refactor(forks): tighten ForkProtocol surface (Stage 1 of #686)

packages/testing/src/consensus_testing/genesis.py

@@ -4,13 +4,16 @@
 from lean_spec.forks.lstar.containers.slot import Slot
 from lean_spec.forks.lstar.containers.state import State, Validators
 from lean_spec.forks.lstar.containers.validator import Validator, ValidatorIndex
+from lean_spec.forks.lstar.spec import LstarSpec
 from lean_spec.forks.protocol import ForkProtocol
 from lean_spec.subspecs.ssz.hash import hash_tree_root
 from lean_spec.types import Bytes52, Uint64
 
 from .keys import XmssKeyManager
 
 _DEFAULT_GENESIS_TIME = Uint64(0)
+_DEFAULT_FORK: ForkProtocol = LstarSpec()
+"""Stateless fork instance used when callers do not pass one explicitly."""
 
 
 def _build_validators(num_validators: int) -> Validators:
@@ -39,34 +42,30 @@ def _build_validators(num_validators: int) -> Validators:
 
 
 def generate_pre_state(
+    fork: ForkProtocol = _DEFAULT_FORK,
     genesis_time: Uint64 = _DEFAULT_GENESIS_TIME,
     num_validators: int = 4,
-    fork: ForkProtocol | None = None,
 ) -> State:
     """Generate a default pre-state for consensus tests.
 
-    A fork argument routes through that fork's own genesis builder.
-    Later forks adding extra state fields stay correct this way.
-
     Args:
+        fork: Fork dispatching genesis construction.
         genesis_time: The genesis timestamp.
         num_validators: Number of validators to include.
-        fork: Optional fork dispatching genesis. Defaults to the base builder.
 
     Returns:
         A properly initialized consensus state.
     """
     validators = _build_validators(num_validators)
-
-    if fork is not None:
-        return fork.generate_genesis(genesis_time=genesis_time, validators=validators)
-
-    return State.generate_genesis(genesis_time=genesis_time, validators=validators)
+    state = fork.generate_genesis(genesis_time=genesis_time, validators=validators)
+    assert isinstance(state, State)
+    return state
 
 
 def build_anchor(
     num_validators: int,
     anchor_slot: Slot,
+    fork: ForkProtocol = _DEFAULT_FORK,
     genesis_time: Uint64 = _DEFAULT_GENESIS_TIME,
 ) -> tuple[State, Block]:
     """Build a consistent non-genesis anchor by advancing the genesis state.
@@ -82,6 +81,7 @@ def build_anchor(
     - State latest_block_header matches the anchor block header (without state_root).
 
     Args:
+        fork: Fork dispatching genesis construction.
         num_validators: Size of the validator set in the anchor state.
         anchor_slot: Slot at which the anchor block lives. Must be > 0.
         genesis_time: Genesis timestamp for the underlying pre-state.
@@ -99,7 +99,7 @@ def build_anchor(
             "For a genesis anchor use generate_pre_state instead."
         )
 
-    state = generate_pre_state(genesis_time=genesis_time, num_validators=num_validators)
+    state = generate_pre_state(fork=fork, genesis_time=genesis_time, num_validators=num_validators)
 
     # Reconstruct the genesis block from the state's latest header.
     # The genesis block is fully determined by the genesis state.

packages/testing/src/consensus_testing/test_fixtures/api_endpoint.py

@@ -8,6 +8,7 @@
 from lean_spec.forks.lstar.containers.block import Block
 from lean_spec.forks.lstar.containers.block.types import AggregatedAttestations
 from lean_spec.forks.lstar.containers.state import State
+from lean_spec.forks.lstar.spec import LstarSpec
 from lean_spec.subspecs.ssz.hash import hash_tree_root
 from lean_spec.types import Bytes32, Uint64
 
@@ -45,8 +46,11 @@ def _build_store(num_validators: int, genesis_time: int, anchor_slot: int = 0) -
     endpoint responses whose shape depends on post-genesis slot numbers,
     historical roots, and multi-node fork-choice trees.
     """
+    fork = LstarSpec()
     if anchor_slot == 0:
-        state = generate_pre_state(genesis_time=Uint64(genesis_time), num_validators=num_validators)
+        state = generate_pre_state(
+            fork=fork, genesis_time=Uint64(genesis_time), num_validators=num_validators
+        )
         block = _make_genesis_block(state)
         # No validator identity — fixture only reads store data, never signs.
         return Store.from_anchor(state, block, validator_id=None)
@@ -55,6 +59,7 @@ def _build_store(num_validators: int, genesis_time: int, anchor_slot: int = 0) -
     # The returned pair (state, block) is internally consistent with the
     # historical chain the fixture wants to present to the endpoint.
     state, block = build_anchor(
+        fork=fork,
         num_validators=num_validators,
         anchor_slot=Slot(anchor_slot),
         genesis_time=Uint64(genesis_time),

packages/testing/src/consensus_testing/test_fixtures/sync.py

@@ -8,6 +8,7 @@
 from typing import Any, ClassVar
 
 from lean_spec.forks.lstar.containers.slot import Slot
+from lean_spec.forks.lstar.spec import LstarSpec
 from lean_spec.subspecs.sync.checkpoint_sync import verify_checkpoint_state
 from lean_spec.types import Uint64
 
@@ -74,10 +75,14 @@ def _make_verify_checkpoint(self) -> dict[str, Any]:
         """
         num_validators = int(self.input["numValidators"])
         anchor_slot = int(self.input.get("anchorSlot", 0))
+        fork = LstarSpec()
         if anchor_slot == 0:
-            state = generate_pre_state(genesis_time=Uint64(0), num_validators=num_validators)
+            state = generate_pre_state(
+                fork=fork, genesis_time=Uint64(0), num_validators=num_validators
+            )
         else:
             state, _ = build_anchor(
+                fork=fork,
                 num_validators=num_validators,
                 anchor_slot=Slot(anchor_slot),
                 genesis_time=Uint64(0),

packages/testing/src/framework/pytest_plugins/filler.py

@@ -422,7 +422,7 @@ def test_case_description(request: pytest.FixtureRequest) -> str:
 
 
 @pytest.fixture(scope="function")
-def pre(request: pytest.FixtureRequest) -> Any:
+def pre(request: pytest.FixtureRequest, fork: Any) -> Any:
     """
     Default pre-state (layer-specific).
 
@@ -438,11 +438,12 @@ def pre(request: pytest.FixtureRequest) -> Any:
         )
 
     layer_module = request.config.layer_module  # type: ignore[attr-defined]
+    spec = fork.spec_class()()
 
     if hasattr(request, "param"):
-        return layer_module.generate_pre_state(**request.param)
+        return layer_module.generate_pre_state(fork=spec, **request.param)
 
-    return layer_module.generate_pre_state()
+    return layer_module.generate_pre_state(fork=spec)
 
 
 def base_spec_filler_parametrizer(fixture_class: Any) -> Any:

src/lean_spec/__main__.py

@@ -32,8 +32,9 @@
 import sys
 import time
 from pathlib import Path
+from typing import cast
 
-from lean_spec.forks import DEFAULT_REGISTRY, ForkProtocol, State
+from lean_spec.forks import DEFAULT_REGISTRY, ForkProtocol, State, Store
 from lean_spec.forks.lstar.containers import Block, BlockBody, Checkpoint
 from lean_spec.forks.lstar.containers.block.types import AggregatedAttestations
 from lean_spec.forks.lstar.containers.slot import Slot
@@ -197,7 +198,7 @@ def _init_from_genesis(
         network=event_source.reqresp_client,
         fork=fork,
         validator_registry=validator_registry,
-        network_name=fork.NETWORK_NAME,
+        network_name=fork.GOSSIP_DIGEST,
         is_aggregator=is_aggregator,
         aggregate_subnet_ids=aggregate_subnet_ids,
         api_config=ApiServerConfig(port=api_port) if api_port is not None else None,
@@ -288,7 +289,7 @@ async def _init_from_checkpoint(
         # The store treats this as the new "genesis" for fork choice purposes.
         # All blocks before the checkpoint are effectively pruned.
         validator_id = validator_registry.primary_index() if validator_registry else None
-        store = fork.create_store(state, anchor_block, validator_id)
+        store = cast(Store, fork.create_store(state, anchor_block, validator_id))
         logger.info(
             "Initialized from checkpoint at slot %d (finalized=%s)",
             state.slot,
@@ -313,7 +314,7 @@ async def _init_from_checkpoint(
             network=event_source.reqresp_client,
             fork=fork,
             validator_registry=validator_registry,
-            network_name=fork.NETWORK_NAME,
+            network_name=fork.GOSSIP_DIGEST,
             is_aggregator=is_aggregator,
             aggregate_subnet_ids=aggregate_subnet_ids,
             api_config=ApiServerConfig(port=api_port) if api_port is not None else None,
@@ -503,14 +504,14 @@ async def run_node(
     #
     # Without this, the event source defaults to "0x00000000" and rejects
     # all messages from other clients that use "devnet0".
-    event_source.set_network_name(fork.NETWORK_NAME)
+    event_source.set_network_name(fork.GOSSIP_DIGEST)
 
     # Subscribe to gossip topics.
     #
     # We subscribe before connecting to bootnodes so that when
     # we establish connections, we can immediately announce our
     # subscriptions to peers.
-    block_topic = GossipTopic.block(fork.NETWORK_NAME).to_topic_id()
+    block_topic = GossipTopic.block(fork.GOSSIP_DIGEST).to_topic_id()
     event_source.subscribe_gossip_topic(block_topic)
 
     # Determine attestation subnets to subscribe to.
@@ -541,7 +542,7 @@ async def run_node(
 
     for subnet_id in subscription_subnets:
         attestation_subnet_topic = GossipTopic.attestation_subnet(
-            fork.NETWORK_NAME, subnet_id
+            fork.GOSSIP_DIGEST, subnet_id
         ).to_topic_id()
         event_source.subscribe_gossip_topic(attestation_subnet_topic)
         logger.info("Subscribed to attestation subnet %d", subnet_id)

src/lean_spec/forks/lstar/spec.py

@@ -4,7 +4,7 @@
 
 from lean_spec.forks.lstar.containers.block import Block
 
-from ..protocol import ForkProtocol
+from ..protocol import ForkProtocol, SpecStateType
 from .containers.state import State
 from .store import Store
 
@@ -14,18 +14,19 @@ class LstarSpec(ForkProtocol):
 
     NAME: ClassVar[str] = "lstar"
     VERSION: ClassVar[int] = 4
-    NETWORK_NAME: ClassVar[str] = "devnet0"
+    GOSSIP_DIGEST: ClassVar[str] = "devnet0"
 
     previous: ClassVar[type[ForkProtocol] | None] = None
 
-    state_class: ClassVar[type[State]] = State
-    block_class: ClassVar[type[Block]] = Block
-    store_class: ClassVar[type[Store]] = Store
+    state_class: type[State] = State
+    block_class: type[Block] = Block
+    store_class: type[Store] = Store
 
-    def upgrade_state(self, state: State) -> State:
+    def upgrade_state(self, state: SpecStateType) -> State:
         """
         Lstar is the root fork: there is no predecessor, so no migration.
 
         Returns the input state unchanged.
         """
+        assert isinstance(state, State)
         return state

src/lean_spec/forks/protocol.py

@@ -7,21 +7,34 @@
 from abc import ABC, abstractmethod
 from typing import Any, ClassVar, Protocol, Self
 
+from lean_spec.types import Bytes32, SSZList, Uint64
+
 
 class SpecStateType(Protocol):
     """Structural contract: any fork's State container class exposes genesis."""
 
     @classmethod
-    def generate_genesis(cls, genesis_time: Any, validators: Any) -> Self:
+    def generate_genesis(cls, genesis_time: Uint64, validators: SSZList[Any]) -> Self:
         """Construct the fork's genesis state."""
         ...
 
 
+class SpecBlockType(Protocol):
+    """Structural contract: any fork's Block container class."""
+
+
 class SpecStoreType(Protocol):
     """Structural contract: any fork's Store class exposes anchor construction."""
 
+    head: Bytes32
+
     @classmethod
-    def from_anchor(cls, state: Any, anchor_block: Any, validator_id: Any) -> Self:
+    def from_anchor(
+        cls,
+        state: SpecStateType,
+        anchor_block: SpecBlockType,
+        validator_id: Uint64 | None,
+    ) -> Self:
         """Construct a forkchoice store anchored at the given state/block."""
         ...
 
@@ -35,7 +48,7 @@ class ForkProtocol(ABC):
     VERSION: ClassVar[int]
     """Strictly monotonic version. Used to order forks in the registry."""
 
-    NETWORK_NAME: ClassVar[str]
+    GOSSIP_DIGEST: ClassVar[str]
     """
     Fork identifier embedded in gossipsub topic names.
 
@@ -51,25 +64,30 @@ class ForkProtocol(ABC):
     and that upgrade_state can traverse for cross-fork state migrations.
     """
 
-    state_class: ClassVar[type[SpecStateType]]
+    state_class: type[SpecStateType]
     """Concrete State container class owned by this fork."""
 
-    block_class: ClassVar[type]
+    block_class: type[SpecBlockType]
     """Concrete Block container class owned by this fork."""
 
-    store_class: ClassVar[type[SpecStoreType]]
+    store_class: type[SpecStoreType]
     """Concrete forkchoice Store class owned by this fork."""
 
-    def generate_genesis(self, genesis_time: Any, validators: Any) -> Any:
+    def generate_genesis(self, genesis_time: Uint64, validators: SSZList[Any]) -> SpecStateType:
         """Construct a genesis state using this fork's State class."""
         return self.state_class.generate_genesis(genesis_time, validators)
 
-    def create_store(self, state: Any, anchor_block: Any, validator_id: Any) -> Any:
+    def create_store(
+        self,
+        state: SpecStateType,
+        anchor_block: SpecBlockType,
+        validator_id: Uint64 | None,
+    ) -> SpecStoreType:
         """Construct a forkchoice store anchored at the given state and block."""
         return self.store_class.from_anchor(state, anchor_block, validator_id)
 
     @abstractmethod
-    def upgrade_state(self, state: Any) -> Any:
+    def upgrade_state(self, state: SpecStateType) -> SpecStateType:
         """
         Migrate state from the previous fork's shape into this fork's shape.
 

src/lean_spec/subspecs/node/node.py

@@ -17,9 +17,9 @@
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Final
+from typing import Final, cast
 
-from lean_spec.forks import ForkProtocol, Store
+from lean_spec.forks import ForkProtocol, State, Store
 from lean_spec.forks.lstar.containers import Block, BlockBody, SignedBlock
 from lean_spec.forks.lstar.containers.attestation import SignedAttestation
 from lean_spec.forks.lstar.containers.block.types import AggregatedAttestations
@@ -212,7 +212,7 @@ def from_genesis(cls, config: NodeConfig) -> Node:
             # Generate genesis state from validators.
             #
             # Includes initial checkpoints, validator registry, and config.
-            state = fork.generate_genesis(config.genesis_time, config.validators)
+            state = cast(State, fork.generate_genesis(config.genesis_time, config.validators))
 
             # Create genesis block.
             #
@@ -229,7 +229,7 @@ def from_genesis(cls, config: NodeConfig) -> Node:
             # Initialize forkchoice store.
             #
             # Genesis block is both justified and finalized.
-            store = fork.create_store(state, block, validator_id)
+            store = cast(Store, fork.create_store(state, block, validator_id))
 
             # Persist genesis to database if available.
             #

tests/lean_spec/forks/test_fork_protocol.py

@@ -59,9 +59,9 @@ def test_identity(self) -> None:
         assert LstarSpec.NAME == "lstar"
         assert LstarSpec.VERSION == 4
 
-    def test_network_name(self) -> None:
-        """LstarSpec carries the gossipsub network name as fork metadata."""
-        assert LstarSpec.NETWORK_NAME == "devnet0"
+    def test_gossip_digest(self) -> None:
+        """LstarSpec carries the gossipsub fork digest as fork metadata."""
+        assert LstarSpec.GOSSIP_DIGEST == "devnet0"
 
     def test_previous_is_none(self) -> None:
         """LstarSpec is the root of the upgrade chain."""
@@ -81,6 +81,7 @@ def test_generate_genesis(self) -> None:
         fork = LstarSpec()
         validators = make_validators(4)
         state = fork.generate_genesis(Uint64(0), validators)
+        assert isinstance(state, State)
         assert state.slot == Slot(0)
         assert len(state.validators) == 4