polyc-state-connect 2026.10.2

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::{
    error::StateError,
    page::Positioned as _,
    query_audit::{ListQueryAuditHistory, QueryAuditHistory, QueryAuditHistoryEntry},
    revision::{JournalPosition, PartitionIncarnation},
};

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

/// State's narrow query-audit history service.
pub struct QueryAuditHistorySvc {
    authority: Arc<dyn QueryAuditHistory>,
    draining: Arc<AtomicBool>,
    binding: Arc<AudienceBinding>,
}

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

    /// Registers the two history routes.
    #[must_use]
    pub fn register_on(self, router: Router) -> Router {
        use pb::StateQueryAuditHistoryServiceExt as _;
        Arc::new(self).register(router)
    }

    fn admit(
        &self,
        ctx: &RequestContext,
        value: impl Into<Option<pb::CallContext>>,
        method: &'static str,
    ) -> Result<
        (
            polyc_state::context::CallContext,
            tracing::Span,
            std::time::Duration,
        ),
        ConnectError,
    > {
        check_not_draining(self.draining.load(Ordering::Relaxed))?;
        let declared: DeclaredCall =
            declared_call(value).map_err(|error| to_connect_error(&error))?;
        let family = polyc_state::query_audit::family();
        check_call_context_version(declared.version).map_err(|error| to_connect_error(&error))?;
        let peer = PeerIdentity::from_verified_leaf(
            ctx.peer_certs().and_then(<[_]>::first).map(|leaf| &**leaf),
        );
        check_audience_binding(
            &peer,
            &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),
            slot_wait(ctx.time_remaining(), &declared),
        ))
    }
}

/// Maps a query-audit refusal onto a transport refusal.
fn audit_error(error: polyc_state::query_audit::QueryAuditError) -> StateError {
    match error {
        polyc_state::query_audit::QueryAuditError::State(state) => state,
        polyc_state::query_audit::QueryAuditError::NoRecordedIntent { query } => {
            StateError::Malformed {
                field: "history".to_owned(),
                reason: format!("the authority answered about {query}, which was not asked"),
            }
        }
    }
}

/// Frames one entry for the wire, without its signature material.
fn entry_message(entry: &QueryAuditHistoryEntry) -> pb::QueryAuditHistoryEntry {
    use polyc_proto::proto::polychrome::state::v1::query_audit_history_entry::Phase;
    let receipt = entry.receipt();
    pb::QueryAuditHistoryEntry {
        ordinal: entry.position().get(),
        command_id: receipt.command_id().as_str().to_owned(),
        receipt_digest: receipt.digest().as_bytes().to_vec(),
        phase: Some(match entry {
            QueryAuditHistoryEntry::Intent { intent, .. } => {
                Phase::Intent(Box::new(pb::QueryAuditIntent::from(Kernel(intent))))
            }
            QueryAuditHistoryEntry::Completion {
                query,
                namespace,
                intent_ordinal,
                completion,
                ..
            } => Phase::Completion(Box::new(pb::QueryAuditHistoryCompletion {
                query: query.as_str().to_owned(),
                namespace: namespace.as_str().to_owned(),
                intent_ordinal: intent_ordinal.get(),
                completion: buffa::MessageField::some(pb::QueryCompletion::from(Kernel(
                    completion,
                ))),
                __buffa_unknown_fields: buffa::UnknownFields::default(),
            })),
        }),
        __buffa_unknown_fields: buffa::UnknownFields::default(),
    }
}

#[allow(refining_impl_trait)]
impl pb::StateQueryAuditHistoryService for QueryAuditHistorySvc {
    async fn get_source(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::GetQueryAuditSourceRequest>,
    ) -> ServiceResult<pb::GetQueryAuditSourceReply> {
        let pb::GetQueryAuditSourceRequest {
            context,
            __buffa_unknown_fields: _,
        } = request.to_owned_message();
        let (context, span, wait) = self.admit(&ctx, context, "GetQueryAuditSource")?;
        let authority = Arc::clone(&self.authority);
        let head = blocking::read(span, polyc_state::query_audit::family(), wait, move || {
            authority.source_head(&context).map_err(audit_error)
        })
        .await?;
        Response::ok(pb::GetQueryAuditSourceReply {
            source: buffa::MessageField::some(pb::AuditSource {
                partition: head.source().partition().as_str().to_owned(),
                incarnation: head.source().incarnation().as_bytes().to_vec(),
                __buffa_unknown_fields: buffa::UnknownFields::default(),
            }),
            head: head.head().get(),
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }

    async fn list_history(
        &self,
        ctx: RequestContext,
        request: ServiceRequest<'_, pb::ListQueryAuditHistoryRequest>,
    ) -> ServiceResult<pb::ListQueryAuditHistoryReply> {
        let pb::ListQueryAuditHistoryRequest {
            context,
            after,
            limit,
            expected_lineage,
            __buffa_unknown_fields: _,
        } = request.to_owned_message();
        let expected_lineage = expected_lineage
            .map(|bytes| {
                <[u8; PartitionIncarnation::LEN]>::try_from(bytes.as_slice()).map_err(|_| {
                    to_connect_error(&StateError::Malformed {
                        field: "expected_lineage".to_owned(),
                        reason: "a lineage is exactly 32 bytes".to_owned(),
                    })
                })
            })
            .transpose()?
            .map(PartitionIncarnation::from_bytes);
        let (context, span, wait) = self.admit(&ctx, context, "ListQueryAuditHistory")?;
        // A resume names the lineage that issued its cursor. Serving one that
        // does not would hand a caller the new log's entries as a continuation
        // of a cursor the old log issued, which is the loss the lineage exists
        // to refuse. A first read binds nothing and learns its lineage here.
        let requested = match (after, expected_lineage) {
            (None, _) => ListQueryAuditHistory::first(limit),
            (Some(after), Some(lineage)) => {
                ListQueryAuditHistory::after(JournalPosition::new(after), lineage, limit)
            }
            (Some(_), None) => {
                return Err(to_connect_error(&StateError::Malformed {
                    field: "expected_lineage".to_owned(),
                    reason: "a resume names the lineage that issued its cursor".to_owned(),
                }));
            }
        };
        let authority = Arc::clone(&self.authority);
        // The watermark comes from the page, so it names the instant the page
        // was read at. The source is a second read under a second acquisition
        // of the authority's lock. That is safe only because the lineage is
        // minted once when the log is created and cannot change while this
        // process holds the authority open: a value that could change between
        // the two reads would have to be carried out of the first.
        let page = blocking::read(span, polyc_state::query_audit::family(), wait, move || {
            let page = authority
                .history(requested, &context)
                .map_err(audit_error)?;
            let head = authority.source_head(&context).map_err(audit_error)?;
            Ok::<_, StateError>((page, head))
        })
        .await?;
        let (page, head) = page;
        Response::ok(pb::ListQueryAuditHistoryReply {
            entries: page.records().iter().map(entry_message).collect(),
            next: page
                .next_cursor()
                .map_or_else(buffa::MessageField::default, |cursor| {
                    buffa::MessageField::some(pb::Cursor::from(Kernel(cursor)))
                }),
            completeness: pb::PageCompleteness::from(Kernel(page.completeness())).into(),
            consistency: pb::Consistency::from(Kernel(page.consistency())).into(),
            watermark: page
                .watermark()
                .map_or(0, polyc_state::consistency::Watermark::get),
            source: buffa::MessageField::some(pb::AuditSource {
                partition: head.source().partition().as_str().to_owned(),
                incarnation: head.source().incarnation().as_bytes().to_vec(),
                __buffa_unknown_fields: buffa::UnknownFields::default(),
            }),
            __buffa_unknown_fields: buffa::UnknownFields::default(),
        })
    }
}