polyc-state-connect 2026.9.0

//! Server adapter for durable agent tasks.

use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};

use connectrpc::{ConnectError, RequestContext, Response, Router, ServiceRequest, ServiceResult};
use polyc_proto::proto::polychrome::state::v1 as pb;
use polyc_state::{
    context::CallContext,
    error::StateError,
    id::{CommandId, NamespaceId},
    tasks::{ContextId, EdgeId, TaskDirectory, TaskId, TaskRead, TaskWrite},
    versioned::{VersionedRead, VersionedTransact},
};

use crate::{
    admission::{
        AudienceBinding, PeerIdentity, check_audience_binding, check_call_context_version,
        check_not_draining, check_transport_deadline, state_audience,
    },
    error::to_connect_error,
    trace::adopt_caller_trace,
    wire::{DeclaredCall, Kernel, declared_call},
};

use super::wire::{command_from_wire, index_to_wire, record_to_wire};

/// Complete, capability-specific task authority port.
pub trait TaskAuthority: TaskRead + TaskWrite {
    /// Namespace this authority owns.
    fn namespace(&self) -> &NamespaceId;
}

impl<V> TaskAuthority for TaskDirectory<V>
where
    V: VersionedRead + VersionedTransact,
{
    fn namespace(&self) -> &NamespaceId {
        self.namespace()
    }
}

/// State's agent-task service.
pub struct TaskSvc {
    authority: Arc<dyn TaskAuthority>,
    draining: Arc<AtomicBool>,
    binding: Arc<AudienceBinding>,
}

impl TaskSvc {
    /// Builds the service.
    #[must_use]
    pub const fn new(
        authority: Arc<dyn TaskAuthority>,
        draining: Arc<AtomicBool>,
        binding: Arc<AudienceBinding>,
    ) -> Self {
        Self {
            authority,
            draining,
            binding,
        }
    }

    /// Registers the capability-specific route family.
    #[must_use]
    pub fn register_on(self, router: Router) -> Router {
        use pb::StateTaskServiceExt as _;
        Arc::new(self).register(router)
    }

    fn peer(ctx: &RequestContext) -> PeerIdentity {
        PeerIdentity::from_verified_leaf(
            ctx.peer_certs().and_then(<[_]>::first).map(|leaf| &**leaf),
        )
    }

    fn admit(
        &self,
        ctx: &RequestContext,
        context: impl Into<Option<pb::CallContext>>,
        method: &'static str,
    ) -> Result<(CallContext, tracing::Span), ConnectError> {
        check_not_draining(self.draining.load(Ordering::Relaxed))?;
        let declared: DeclaredCall =
            declared_call(context).map_err(|error| to_connect_error(&error))?;
        let family = polyc_state::tasks::family();
        check_call_context_version(declared.version).map_err(|error| to_connect_error(&error))?;
        check_audience_binding(
            &Self::peer(ctx),
            &declared.audience,
            &state_audience(),
            &self.binding,
            &family,
        )
        .map_err(|error| to_connect_error(&error))?;
        check_transport_deadline(ctx.time_remaining(), &family)
            .map_err(|error| to_connect_error(&error))?;
        Ok((
            declared.origin_relative_context(),
            adopt_caller_trace(ctx.headers(), method),
        ))
    }

    fn namespace(&self, value: &str) -> Result<(), ConnectError> {
        if value == self.authority.namespace().as_str() {
            Ok(())
        } else {
            Err(to_connect_error(&StateError::Denied {
                family: polyc_state::tasks::family(),
            }))
        }
    }
}

#[allow(refining_impl_trait)]
impl pb::StateTaskService for TaskSvc {
    async fn transact(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::TransactStateTaskRequest>,
    ) -> ServiceResult<pb::TransactStateTaskReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "TransactStateTask")?;
        let _entered = span.enter();
        let metadata = message
            .metadata
            .into_option()
            .ok_or_else(|| malformed_connect("metadata", "a task command carries metadata"))?;
        self.namespace(&metadata.namespace)?;
        let operation = message.operation.into_option().ok_or_else(|| {
            malformed_connect("operation", "a task command carries one operation")
        })?;
        let receipt = self
            .authority
            .transact(
                command_from_wire(metadata, operation).map_err(|error| to_connect_error(&error))?,
                &context,
            )
            .map_err(|error| to_connect_error(&error))?;
        Response::ok(pb::TransactStateTaskReply {
            receipt: buffa::MessageField::some(pb::Receipt::from(Kernel(&receipt))),
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }

    async fn get_task(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::GetStateTaskRequest>,
    ) -> ServiceResult<pb::GetStateTaskReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "GetStateTask")?;
        let _entered = span.enter();
        self.namespace(&message.namespace)?;
        let fact = self
            .authority
            .task(
                &EdgeId::new(message.owner_edge_id),
                &TaskId::new(message.task_id),
                &context,
            )
            .map_err(|error| to_connect_error(&error))?;
        let snapshot_revision = fact.snapshot_revision().get();
        let entry_revision = fact
            .entry_revision()
            .map(polyc_state::revision::Revision::get);
        Response::ok(pb::GetStateTaskReply {
            task: fact
                .value()
                .as_ref()
                .map_or_else(buffa::MessageField::none, |record| {
                    buffa::MessageField::some(record_to_wire(record))
                }),
            snapshot_revision,
            entry_revision,
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }

    async fn get_context_index(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::GetStateTaskContextIndexRequest>,
    ) -> ServiceResult<pb::GetStateTaskContextIndexReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "GetStateTaskContextIndex")?;
        let _entered = span.enter();
        self.namespace(&message.namespace)?;
        let fact = self
            .authority
            .context_index(
                &EdgeId::new(message.owner_edge_id),
                &ContextId::new(message.context_id),
                &context,
            )
            .map_err(|error| to_connect_error(&error))?;
        Response::ok(pb::GetStateTaskContextIndexReply {
            index: buffa::MessageField::some(index_to_wire(fact.value())),
            snapshot_revision: fact.snapshot_revision().get(),
            entry_revision: fact
                .entry_revision()
                .map(polyc_state::revision::Revision::get),
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }

    async fn list_tasks(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::ListStateTasksRequest>,
    ) -> ServiceResult<pb::ListStateTasksReply> {
        let message = request.to_owned_message();
        let (context, span) = self.admit(&ctx, message.context, "ListStateTasks")?;
        let _entered = span.enter();
        self.namespace(&message.namespace)?;
        let after = message.after.map(TaskId::new);
        let fact = self
            .authority
            .list_by_context(
                &EdgeId::new(message.owner_edge_id),
                &ContextId::new(message.context_id),
                after.as_ref(),
                usize::try_from(message.page_size).unwrap_or(usize::MAX),
                &context,
            )
            .map_err(|error| to_connect_error(&error))?;
        let snapshot_revision = fact.snapshot_revision().get();
        let entry_revision = fact
            .entry_revision()
            .map(polyc_state::revision::Revision::get);
        let page = fact.into_value();
        Response::ok(pb::ListStateTasksReply {
            tasks: page.tasks().iter().map(record_to_wire).collect(),
            next_after: page.next_after().map(|id| id.as_str().to_owned()),
            snapshot_revision,
            entry_revision,
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }

    async fn get_receipt(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::GetStateTaskReceiptRequest>,
    ) -> ServiceResult<pb::GetStateTaskReceiptReply> {
        let message = request.to_owned_message();
        let (_context, span) = self.admit(&ctx, message.context, "GetStateTaskReceipt")?;
        let _entered = span.enter();
        self.namespace(&message.namespace)?;
        let receipt = self
            .authority
            .committed_receipt(
                &NamespaceId::new(message.namespace),
                &CommandId::new(message.command_id),
            )
            .map_err(|error| to_connect_error(&error))?;
        Response::ok(pb::GetStateTaskReceiptReply {
            receipt: receipt
                .as_ref()
                .map_or_else(buffa::MessageField::none, |value| {
                    buffa::MessageField::some(pb::Receipt::from(Kernel(value)))
                }),
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }
}

fn malformed_connect(field: &str, reason: &str) -> ConnectError {
    to_connect_error(&StateError::Malformed {
        field: field.into(),
        reason: reason.into(),
    })
}