diff --git a/src/lean_spec/subspecs/api/endpoints/__init__.py b/src/lean_spec/subspecs/api/endpoints/__init__.py index e2c0f12f..05654044 100644 --- a/src/lean_spec/subspecs/api/endpoints/__init__.py +++ b/src/lean_spec/subspecs/api/endpoints/__init__.py @@ -1,9 +1,10 @@ """API endpoint specifications.""" -from . import checkpoints, health, states +from . import checkpoints, fork_choice, health, states __all__ = [ "checkpoints", + "fork_choice", "health", "states", ] diff --git a/src/lean_spec/subspecs/api/endpoints/fork_choice.py b/src/lean_spec/subspecs/api/endpoints/fork_choice.py new file mode 100644 index 00000000..7c9a5c74 --- /dev/null +++ b/src/lean_spec/subspecs/api/endpoints/fork_choice.py @@ -0,0 +1,74 @@ +"""Fork choice endpoint handler.""" + +from __future__ import annotations + +import json + +from aiohttp import web + + +async def handle(request: web.Request) -> web.Response: + """ + Handle fork choice tree request. + + Returns the fork choice tree snapshot: blocks with weights, head, + checkpoints, safe target, and validator count. + + Response: JSON object with fields: + - nodes (array): Blocks in the tree, each with root, slot, parent_root, + proposer_index, and weight. + - head (string): Current head block root as 0x-prefixed hex. + - justified (object): Latest justified checkpoint (slot, root). + - finalized (object): Latest finalized checkpoint (slot, root). + - safe_target (string): Safe target block root as 0x-prefixed hex. + - validator_count (integer): Number of validators in head state. + + Status Codes: + 200 OK: Fork choice tree returned successfully. + 503 Service Unavailable: Store not initialized. + """ + store_getter = request.app.get("store_getter") + store = store_getter() if store_getter else None + + if store is None: + raise web.HTTPServiceUnavailable(reason="Store not initialized") + + finalized_slot = store.latest_finalized.slot + weights = store.compute_block_weights() + + nodes = [] + for root, block in store.blocks.items(): + if block.slot < finalized_slot: + continue + nodes.append( + { + "root": "0x" + root.hex(), + "slot": int(block.slot), + "parent_root": "0x" + block.parent_root.hex(), + "proposer_index": int(block.proposer_index), + "weight": weights.get(root, 0), + } + ) + + head_state = store.states.get(store.head) + validator_count = len(head_state.validators) if head_state is not None else 0 + + response = { + "nodes": nodes, + "head": "0x" + store.head.hex(), + "justified": { + "slot": int(store.latest_justified.slot), + "root": "0x" + store.latest_justified.root.hex(), + }, + "finalized": { + "slot": int(store.latest_finalized.slot), + "root": "0x" + store.latest_finalized.root.hex(), + }, + "safe_target": "0x" + store.safe_target.hex(), + "validator_count": validator_count, + } + + return web.Response( + body=json.dumps(response), + content_type="application/json", + ) diff --git a/src/lean_spec/subspecs/api/routes.py b/src/lean_spec/subspecs/api/routes.py index ed08af6b..a06eefa5 100644 --- a/src/lean_spec/subspecs/api/routes.py +++ b/src/lean_spec/subspecs/api/routes.py @@ -6,11 +6,12 @@ from aiohttp import web -from .endpoints import checkpoints, health, states +from .endpoints import checkpoints, fork_choice, health, states ROUTES: dict[str, Callable[[web.Request], Awaitable[web.Response]]] = { "/lean/v0/health": health.handle, "/lean/v0/states/finalized": states.handle_finalized, "/lean/v0/checkpoints/justified": checkpoints.handle_justified, + "/lean/v0/fork_choice": fork_choice.handle, } """All API routes mapped to their handlers.""" diff --git a/src/lean_spec/subspecs/forkchoice/store.py b/src/lean_spec/subspecs/forkchoice/store.py index 58dd573a..3eeb93af 100644 --- a/src/lean_spec/subspecs/forkchoice/store.py +++ b/src/lean_spec/subspecs/forkchoice/store.py @@ -645,6 +645,33 @@ def extract_attestations_from_aggregated_payloads( return attestations + def compute_block_weights(self) -> dict[Bytes32, int]: + """ + Compute attestation-based weight for each block above the finalized slot. + + Walks backward from each validator's latest head vote, incrementing weight + for every ancestor above the finalized slot. + + Returns: + Mapping from block root to accumulated attestation weight. + """ + attestations = self.extract_attestations_from_aggregated_payloads( + self.latest_known_aggregated_payloads + ) + + start_slot = self.latest_finalized.slot + + weights: dict[Bytes32, int] = defaultdict(int) + + for attestation_data in attestations.values(): + current_root = attestation_data.head.root + + while current_root in self.blocks and self.blocks[current_root].slot > start_slot: + weights[current_root] += 1 + current_root = self.blocks[current_root].parent_root + + return dict(weights) + def _compute_lmd_ghost_head( self, start_root: Bytes32, diff --git a/tests/lean_spec/subspecs/api/test_server.py b/tests/lean_spec/subspecs/api/test_server.py index 860ec606..57940860 100644 --- a/tests/lean_spec/subspecs/api/test_server.py +++ b/tests/lean_spec/subspecs/api/test_server.py @@ -102,3 +102,72 @@ async def test_returns_503_when_store_not_initialized(self) -> None: finally: server.stop() await asyncio.sleep(0.1) + + +class TestForkChoiceEndpoint: + """Tests for the /lean/v0/fork_choice endpoint.""" + + async def test_returns_503_when_store_not_initialized(self) -> None: + """Endpoint returns 503 Service Unavailable when store is not set.""" + config = ApiServerConfig(port=15058) + server = ApiServer(config=config) + + await server.start() + + try: + async with httpx.AsyncClient() as client: + response = await client.get("http://127.0.0.1:15058/lean/v0/fork_choice") + + assert response.status_code == 503 + + finally: + server.stop() + await asyncio.sleep(0.1) + + async def test_returns_200_with_initialized_store(self, base_store: Store) -> None: + """Endpoint returns 200 with fork choice tree when store is initialized.""" + config = ApiServerConfig(port=15059) + server = ApiServer(config=config, store_getter=lambda: base_store) + + await server.start() + + try: + async with httpx.AsyncClient() as client: + response = await client.get("http://127.0.0.1:15059/lean/v0/fork_choice") + + assert response.status_code == 200 + assert response.headers["content-type"] == "application/json" + + data = response.json() + + assert set(data.keys()) == { + "nodes", + "head", + "justified", + "finalized", + "safe_target", + "validator_count", + } + + head_root = "0x" + base_store.head.hex() + + assert data["head"] == head_root + assert data["validator_count"] == 3 + assert data["justified"] == { + "slot": int(base_store.latest_justified.slot), + "root": "0x" + base_store.latest_justified.root.hex(), + } + assert data["finalized"] == { + "slot": int(base_store.latest_finalized.slot), + "root": "0x" + base_store.latest_finalized.root.hex(), + } + + assert len(data["nodes"]) == 1 + node = data["nodes"][0] + assert node["root"] == head_root + assert node["slot"] == 0 + assert node["weight"] == 0 + + finally: + server.stop() + await asyncio.sleep(0.1) diff --git a/tests/lean_spec/subspecs/forkchoice/test_compute_block_weights.py b/tests/lean_spec/subspecs/forkchoice/test_compute_block_weights.py new file mode 100644 index 00000000..3233743f --- /dev/null +++ b/tests/lean_spec/subspecs/forkchoice/test_compute_block_weights.py @@ -0,0 +1,134 @@ +"""Tests for Store.compute_block_weights.""" + +from __future__ import annotations + +from lean_spec.subspecs.containers.attestation import AggregationBits, AttestationData +from lean_spec.subspecs.containers.checkpoint import Checkpoint +from lean_spec.subspecs.containers.slot import Slot +from lean_spec.subspecs.containers.validator import ValidatorIndex, ValidatorIndices +from lean_spec.subspecs.forkchoice import Store +from lean_spec.subspecs.ssz.hash import hash_tree_root +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof, SignatureKey +from lean_spec.types.byte_arrays import ByteListMiB +from tests.lean_spec.helpers import make_bytes32, make_signed_block + + +def _make_empty_proof(participants: list[ValidatorIndex]) -> AggregatedSignatureProof: + """Create an aggregated proof with empty proof data for testing.""" + return AggregatedSignatureProof( + participants=AggregationBits.from_validator_indices(ValidatorIndices(data=participants)), + proof_data=ByteListMiB(data=b""), + ) + + +def test_genesis_only_store_returns_empty_weights(base_store: Store) -> None: + """A genesis-only store with no attestations has no block weights.""" + assert base_store.compute_block_weights() == {} + + +def test_linear_chain_weight_accumulates_upward(base_store: Store) -> None: + """Weights walk up from the attested head through all ancestors above finalized slot.""" + genesis_root = base_store.head + + block1 = make_signed_block( + slot=Slot(1), + proposer_index=ValidatorIndex(0), + parent_root=genesis_root, + state_root=make_bytes32(10), + ) + block1_root = hash_tree_root(block1.message.block) + + block2 = make_signed_block( + slot=Slot(2), + proposer_index=ValidatorIndex(1), + parent_root=block1_root, + state_root=make_bytes32(20), + ) + block2_root = hash_tree_root(block2.message.block) + + new_blocks = dict(base_store.blocks) + new_blocks[block1_root] = block1.message.block + new_blocks[block2_root] = block2.message.block + + new_states = dict(base_store.states) + genesis_state = base_store.states[genesis_root] + new_states[block1_root] = genesis_state + new_states[block2_root] = genesis_state + + att_data = AttestationData( + slot=Slot(2), + head=Checkpoint(root=block2_root, slot=Slot(2)), + target=Checkpoint(root=block2_root, slot=Slot(2)), + source=Checkpoint(root=genesis_root, slot=Slot(0)), + ) + data_root = att_data.data_root_bytes() + + proof = _make_empty_proof([ValidatorIndex(0)]) + aggregated_payloads = { + SignatureKey(ValidatorIndex(0), data_root): [proof], + } + + store = base_store.model_copy( + update={ + "blocks": new_blocks, + "states": new_states, + "head": block2_root, + "latest_known_aggregated_payloads": aggregated_payloads, + "attestation_data_by_root": {data_root: att_data}, + } + ) + + weights = store.compute_block_weights() + + # Validator 0 attests to block2 as head. + # Walking up: block2 (slot 2 > 0) gets +1, block1 (slot 1 > 0) gets +1. + # Genesis (slot 0) is at finalized_slot so it does NOT get weight. + assert weights == {block2_root: 1, block1_root: 1} + + +def test_multiple_attestations_accumulate(base_store: Store) -> None: + """Multiple validators attesting to the same head accumulate weight.""" + genesis_root = base_store.head + + block1 = make_signed_block( + slot=Slot(1), + proposer_index=ValidatorIndex(0), + parent_root=genesis_root, + state_root=make_bytes32(10), + ) + block1_root = hash_tree_root(block1.message.block) + + new_blocks = dict(base_store.blocks) + new_blocks[block1_root] = block1.message.block + + new_states = dict(base_store.states) + new_states[block1_root] = base_store.states[genesis_root] + + att_data = AttestationData( + slot=Slot(1), + head=Checkpoint(root=block1_root, slot=Slot(1)), + target=Checkpoint(root=block1_root, slot=Slot(1)), + source=Checkpoint(root=genesis_root, slot=Slot(0)), + ) + data_root = att_data.data_root_bytes() + + proof = _make_empty_proof([ValidatorIndex(0), ValidatorIndex(1)]) + aggregated_payloads = { + SignatureKey(ValidatorIndex(0), data_root): [proof], + SignatureKey(ValidatorIndex(1), data_root): [proof], + } + + store = base_store.model_copy( + update={ + "blocks": new_blocks, + "states": new_states, + "head": block1_root, + "latest_known_aggregated_payloads": aggregated_payloads, + "attestation_data_by_root": {data_root: att_data}, + } + ) + + weights = store.compute_block_weights() + + # Both validators contribute weight to block1 + assert weights == {block1_root: 2}