Duchy Protocol APIs

import grpc
from wfa.measurement.system.v1alpha import computations_service_pb2
from wfa.measurement.system.v1alpha import computations_service_pb2_grpc

def stream_active_computations(stub, continuation_token=None):
    """
    Stream active computations for this duchy.
    
    This is a long-lived stream that yields computations as they
    become available or are updated.
    """
    request = computations_service_pb2.StreamActiveComputationsRequest(
        continuation_token=continuation_token or ""
    )
    
    try:
        for response in stub.StreamActiveComputations(request):
            computation = response.computation
            continuation_token = response.continuation_token
            
            # Persist token for crash recovery
            save_continuation_token(continuation_token)
            
            # Process computation
            print(f"Received computation: {computation.name}")
            print(f"State: {computation.state}")
            
            # Claim and process work
            if should_process_computation(computation):
                process_computation_work(computation)
                
    except grpc.RpcError as e:
        print(f"Stream error: {e.code()} - {e.details()}")
        # Reconnect with last known continuation token
        stream_active_computations(stub, continuation_token)

def submit_computation_result(
    stub,
    computation_name: str,
    aggregator_cert_name: str,
    encrypted_result: bytes,
    result_public_key: bytes
):
    """
    Submit final encrypted result for a computation.
    
    Args:
        stub: Computations service gRPC stub
        computation_name: Computation resource name
        aggregator_cert_name: Aggregator duchy certificate name
        encrypted_result: Encrypted and signed result bytes
        result_public_key: Measurement consumer's public key
    """
    request = computations_service_pb2.SetComputationResultRequest(
        name=computation_name,
        aggregator_certificate=aggregator_cert_name,
        result_public_key=result_public_key,
        encrypted_result=encrypted_result,
        public_api_version="v2alpha"
    )
    
    computation = stub.SetComputationResult(request)
    print(f"Result submitted for: {computation.name}")
    print(f"New state: {computation.state}")
    return computation

def get_assigned_requisitions(computation, duchy_id):
    """
    Get requisitions assigned to a specific duchy.
    
    Args:
        computation: Computation message
        duchy_id: Identifier of the duchy
    
    Returns:
        List of Requisition messages
    """
    assigned_requisitions = []
    
    for requisition in computation.requisitions:
        # Extract duchy from fulfilling_computation_participant
        if duchy_id in requisition.fulfilling_computation_participant:
            assigned_requisitions.append(requisition)
    
    return assigned_requisitions

class DuchyWorker:
    def __init__(self, duchy_id, stub):
        self.duchy_id = duchy_id
        self.stub = stub
        self.active_computations = {}
    
    def run(self):
        """Main worker loop."""
        continuation_token = self._load_continuation_token()
        
        for response in self.stub.StreamActiveComputations(
            computations_service_pb2.StreamActiveComputationsRequest(
                continuation_token=continuation_token
            )
        ):
            computation = response.computation
            continuation_token = response.continuation_token
            
            # Persist token
            self._save_continuation_token(continuation_token)
            
            # Check if this duchy should process
            if self._is_participant(computation):
                # Claim work if not already processing
                if computation.name not in self.active_computations:
                    self._claim_computation(computation)
                
                # Process computation stage
                self._process_computation_stage(computation)
    
    def _is_participant(self, computation):
        """Check if this duchy is a participant."""
        for participant in computation.computation_participants:
            if participant.duchy_id == self.duchy_id:
                return True
        return False
    
    def _claim_computation(self, computation):
        """Claim computation for processing."""
        self.active_computations[computation.name] = computation
        print(f"Claimed computation: {computation.name}")
    
    def _process_computation_stage(self, computation):
        """Process current stage of computation."""
        # Implementation depends on protocol and current stage
        pass