use std::sync::Arc;
use polyc_query_model::GrantSubject;
use crate::session::{MemorySources, QueryScope};
#[async_trait::async_trait]
pub trait PersonaSource: Send + Sync {
async fn active_persona(
&self,
persona_id: String,
) -> Result<Option<polyc_persona::ActivePersona>, polyc_persona::PersonaError>;
async fn participations(
&self,
persona_id: String,
) -> Result<
Vec<polyc_proto::proto::polychrome::persona::v1::Participation>,
polyc_persona::PersonaError,
>;
async fn participation_scope(
&self,
persona_id: String,
cap: usize,
) -> Result<polyc_persona::ScopeResolution, polyc_persona::PersonaError>;
async fn usage_rollup_index(&self) -> Result<Vec<String>, polyc_persona::PersonaError>;
async fn reference_snapshot(
&self,
persona_id: String,
) -> Result<Option<polyc_persona::PersonaReferenceSnapshot>, polyc_persona::PersonaError>;
}
#[derive(Clone)]
pub struct PersonaSourceHandle(Option<Arc<dyn PersonaSource>>);
impl PersonaSourceHandle {
#[must_use]
pub fn current(source: Arc<dyn PersonaSource>) -> Self {
Self(Some(source))
}
#[must_use]
pub const fn unavailable() -> Self {
Self(None)
}
#[must_use]
pub fn load_full(&self) -> Option<Arc<dyn PersonaSource>> {
self.0.clone()
}
}
pub(crate) type PersonaAccess = PersonaSourceHandle;
#[derive(Debug, Clone)]
pub struct AdminPrincipal {
persona_id: String,
}
impl AdminPrincipal {
pub(crate) const fn minted(persona_id: String) -> Self {
Self { persona_id }
}
#[must_use]
pub fn persona_id(&self) -> &str {
&self.persona_id
}
}
#[derive(Debug, Clone)]
pub struct ConversationGrantPrincipal {
conversation_id: String,
subject: GrantSubject,
memory_sources: Vec<String>,
}
impl ConversationGrantPrincipal {
pub(crate) const fn minted(
conversation_id: String,
subject: GrantSubject,
memory_sources: Vec<String>,
) -> Self {
Self {
conversation_id,
subject,
memory_sources,
}
}
#[must_use]
pub(crate) fn memory_sources(&self) -> &[String] {
&self.memory_sources
}
#[must_use]
pub fn conversation_id(&self) -> &str {
&self.conversation_id
}
#[must_use]
pub const fn subject(&self) -> &GrantSubject {
&self.subject
}
#[must_use]
#[allow(
clippy::missing_const_for_fn,
reason = "GrantSubject::turn_id isn't const either — see its own doc"
)]
pub fn turn_id(&self) -> Option<&str> {
self.subject.turn_id()
}
}
#[derive(Debug, Clone)]
pub struct PersonaPrincipal {
persona_id: String,
}
impl PersonaPrincipal {
pub(crate) const fn minted(persona_id: String) -> Self {
Self { persona_id }
}
#[must_use]
pub fn persona_id(&self) -> &str {
&self.persona_id
}
}
#[derive(Debug, Clone)]
pub struct ControlFleetPrincipal {
purpose: String,
statement_digest: [u8; 32],
}
impl ControlFleetPrincipal {
#[cfg(not(any(test, feature = "test-util")))]
pub(crate) const fn minted(purpose: String, statement_digest: [u8; 32]) -> Self {
Self {
purpose,
statement_digest,
}
}
#[cfg(any(test, feature = "test-util"))]
#[must_use]
pub const fn minted(purpose: String, statement_digest: [u8; 32]) -> Self {
Self {
purpose,
statement_digest,
}
}
#[must_use]
pub fn purpose(&self) -> &str {
&self.purpose
}
#[must_use]
pub const fn statement_digest(&self) -> &[u8; 32] {
&self.statement_digest
}
}
#[derive(Debug, Clone)]
pub struct CompositeTracePrincipal {
trace: String,
memory: Option<String>,
routine: String,
addresses: Option<String>,
}
impl CompositeTracePrincipal {
#[cfg(not(any(test, feature = "test-util")))]
pub(crate) fn minted(trace: &str, memory: &str, routine: &str) -> Self {
Self {
trace: trace.to_owned(),
memory: Some(memory.to_owned()),
routine: routine.to_owned(),
addresses: None,
}
}
#[cfg(any(test, feature = "test-util"))]
#[must_use]
pub fn minted(trace: &str, memory: &str, routine: &str) -> Self {
Self {
trace: trace.to_owned(),
memory: Some(memory.to_owned()),
routine: routine.to_owned(),
addresses: None,
}
}
#[must_use]
pub fn trace_statement_digest(&self) -> &str {
&self.trace
}
#[cfg(not(any(test, feature = "test-util")))]
pub(crate) fn minted_without_memory(trace: &str, routine: &str, addresses: &str) -> Self {
Self {
trace: trace.to_owned(),
memory: None,
routine: routine.to_owned(),
addresses: Some(addresses.to_owned()),
}
}
#[cfg(any(test, feature = "test-util"))]
#[must_use]
pub fn minted_without_memory(trace: &str, routine: &str, addresses: &str) -> Self {
Self {
trace: trace.to_owned(),
memory: None,
routine: routine.to_owned(),
addresses: Some(addresses.to_owned()),
}
}
#[must_use]
pub fn memory_statement_digest(&self) -> Option<&str> {
self.memory.as_deref()
}
#[must_use]
pub fn routine_statement_digest(&self) -> &str {
&self.routine
}
#[must_use]
pub fn address_statement_digest(&self) -> Option<&str> {
self.addresses.as_deref()
}
}
#[derive(Debug, Clone)]
pub struct AdminFleetPrincipal {
admin_persona: String,
session: String,
}
impl AdminFleetPrincipal {
#[cfg(not(any(test, feature = "test-util")))]
pub(crate) const fn minted(admin_persona: String, session: String) -> Self {
Self {
admin_persona,
session,
}
}
#[cfg(any(test, feature = "test-util"))]
#[must_use]
pub const fn minted(admin_persona: String, session: String) -> Self {
Self {
admin_persona,
session,
}
}
#[must_use]
pub fn admin_persona(&self) -> &str {
&self.admin_persona
}
#[must_use]
pub fn session(&self) -> &str {
&self.session
}
}
#[derive(Debug, Clone)]
pub enum Principal {
Admin(AdminPrincipal),
ConversationGrant(ConversationGrantPrincipal),
Persona(PersonaPrincipal),
ControlFleet(ControlFleetPrincipal),
AdminFleet(AdminFleetPrincipal),
}
#[derive(Debug, thiserror::Error)]
pub enum PrincipalError {
#[error("no valid admin session")]
InvalidSession,
#[error("persona store unavailable")]
StoreUnavailable,
#[error("this account is not authorized for fleet-wide queries")]
NotAuthorizedForFleet,
#[error("conversation grant token invalid")]
InvalidGrant,
#[error("conversation grant token expired")]
GrantExpired,
#[error("this capability is standalone Query plane only")]
ControlFleetNotUsableEmbedded,
#[error("this capability is standalone Query plane only")]
AdminFleetNotUsableEmbedded,
#[error("a proposed persona-memory source is outside this grant's scope")]
SourceOutsideScope,
#[error("this grant's subject cannot be scoped to a conversation")]
MalformedConversationGrantSubject,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PrincipalKind {
Admin,
ConversationGrantTurn,
ConversationGrantWebSession,
ConversationGrantPersona,
ConversationGrantCompositeTrace,
ConversationGrantAdminCompositeTrace,
Persona,
ControlFleet,
AdminFleet,
}
#[derive(Clone)]
pub struct Scoping {
scope: QueryScope,
allow_explain: bool,
caller_identity: Option<String>,
conversation_id: Option<String>,
turn_id: Option<String>,
web_session_id: Option<String>,
control_fleet: Option<ControlFleetPrincipal>,
admin_fleet: Option<AdminFleetPrincipal>,
composite_trace: Option<CompositeTracePrincipal>,
principal_kind: PrincipalKind,
}
pub struct ScopingParts {
pub scope: QueryScope,
pub allow_explain: bool,
pub caller_identity: Option<String>,
pub conversation_id: Option<String>,
pub turn_id: Option<String>,
pub web_session_id: Option<String>,
}
impl Scoping {
pub(crate) fn for_conversation(
conversation_id: &str,
subject: &GrantSubject,
) -> Result<Self, PrincipalError> {
let composite_trace = subject
.composite_trace()
.map(|(_, trace, memory, routine)| {
CompositeTracePrincipal::minted(trace, memory, routine)
})
.or_else(|| {
subject
.admin_composite_trace()
.map(|(trace, routine, addresses)| {
CompositeTracePrincipal::minted_without_memory(trace, routine, addresses)
})
});
let principal_kind = match subject {
GrantSubject::Turn(_) => PrincipalKind::ConversationGrantTurn,
GrantSubject::WebSession(_) => PrincipalKind::ConversationGrantWebSession,
GrantSubject::Persona(_) => PrincipalKind::ConversationGrantPersona,
GrantSubject::CompositeTrace { .. } => PrincipalKind::ConversationGrantCompositeTrace,
GrantSubject::AdminCompositeTrace { .. } => {
PrincipalKind::ConversationGrantAdminCompositeTrace
}
GrantSubject::ControlFleet { .. } | GrantSubject::AdminFleet { .. } => {
tracing::error!(
"verify_conversation_grant's own invariant broke: a ControlFleet or \
AdminFleet subject reached Scoping::for_conversation wrapped in a \
Principal::ConversationGrant"
);
return Err(PrincipalError::MalformedConversationGrantSubject);
}
};
Ok(Self {
scope: QueryScope::Conversations {
conversations: vec![conversation_id.to_owned()],
memory: MemorySources::default(),
},
allow_explain: false,
caller_identity: subject.persona_id().map(str::to_owned),
conversation_id: Some(conversation_id.to_owned()),
turn_id: subject.turn_id().map(str::to_owned),
web_session_id: subject.web_session_id().map(str::to_owned),
control_fleet: None,
admin_fleet: None,
composite_trace,
principal_kind,
})
}
pub(crate) fn set_scope(&mut self, scope: QueryScope) {
self.scope = scope;
}
pub(crate) fn admin(persona_id: &str) -> Self {
Self {
scope: QueryScope::Fleet,
allow_explain: true,
caller_identity: Some(persona_id.to_owned()),
conversation_id: None,
turn_id: None,
web_session_id: None,
control_fleet: None,
admin_fleet: None,
composite_trace: None,
principal_kind: PrincipalKind::Admin,
}
}
pub(crate) fn persona(persona_id: &str, conversation_ids: Vec<String>) -> Self {
Self {
scope: QueryScope::Conversations {
conversations: conversation_ids,
memory: MemorySources {
owner: Some(persona_id.to_owned()),
participants: Vec::new(),
},
},
allow_explain: false,
caller_identity: Some(persona_id.to_owned()),
conversation_id: None,
turn_id: None,
web_session_id: None,
control_fleet: None,
admin_fleet: None,
composite_trace: None,
principal_kind: PrincipalKind::Persona,
}
}
pub(crate) const fn from_control_fleet(control_fleet: ControlFleetPrincipal) -> Self {
Self {
scope: QueryScope::Fleet,
allow_explain: false,
caller_identity: None,
conversation_id: None,
turn_id: None,
web_session_id: None,
control_fleet: Some(control_fleet),
admin_fleet: None,
composite_trace: None,
principal_kind: PrincipalKind::ControlFleet,
}
}
pub(crate) const fn from_admin_fleet(
caller_persona_id: String,
admin_fleet: AdminFleetPrincipal,
) -> Self {
Self {
scope: QueryScope::Fleet,
allow_explain: true,
caller_identity: Some(caller_persona_id),
conversation_id: None,
turn_id: None,
web_session_id: None,
control_fleet: None,
admin_fleet: Some(admin_fleet),
composite_trace: None,
principal_kind: PrincipalKind::AdminFleet,
}
}
pub(crate) fn own_rows(persona_id: &str, conversation_ids: Vec<String>) -> Self {
Self {
scope: QueryScope::Conversations {
conversations: conversation_ids,
memory: MemorySources::default(),
},
allow_explain: false,
caller_identity: Some(persona_id.to_owned()),
conversation_id: None,
turn_id: None,
web_session_id: None,
control_fleet: None,
admin_fleet: None,
composite_trace: None,
principal_kind: PrincipalKind::Persona,
}
}
#[must_use]
pub const fn scope(&self) -> &QueryScope {
&self.scope
}
#[must_use]
pub const fn allow_explain(&self) -> bool {
self.allow_explain
}
#[must_use]
pub fn caller_identity(&self) -> Option<&str> {
self.caller_identity.as_deref()
}
#[must_use]
pub fn conversation_id(&self) -> Option<&str> {
self.conversation_id.as_deref()
}
#[must_use]
pub fn turn_id(&self) -> Option<&str> {
self.turn_id.as_deref()
}
#[must_use]
pub fn web_session_id(&self) -> Option<&str> {
self.web_session_id.as_deref()
}
#[must_use]
pub const fn control_fleet(&self) -> Option<&ControlFleetPrincipal> {
self.control_fleet.as_ref()
}
#[must_use]
pub const fn admin_fleet(&self) -> Option<&AdminFleetPrincipal> {
self.admin_fleet.as_ref()
}
#[must_use]
pub const fn composite_trace(&self) -> Option<&CompositeTracePrincipal> {
self.composite_trace.as_ref()
}
#[must_use]
pub const fn principal_kind(&self) -> PrincipalKind {
self.principal_kind
}
#[must_use]
pub fn into_parts(self) -> ScopingParts {
ScopingParts {
scope: self.scope,
allow_explain: self.allow_explain,
caller_identity: self.caller_identity,
conversation_id: self.conversation_id,
turn_id: self.turn_id,
web_session_id: self.web_session_id,
}
}
#[cfg(any(test, feature = "test-util"))]
#[must_use]
#[allow(
clippy::too_many_arguments,
reason = "one field per Scoping field, no grouping reduces this"
)]
pub const fn for_test(
scope: QueryScope,
allow_explain: bool,
caller_identity: Option<String>,
conversation_id: Option<String>,
turn_id: Option<String>,
web_session_id: Option<String>,
control_fleet: Option<ControlFleetPrincipal>,
admin_fleet: Option<AdminFleetPrincipal>,
composite_trace: Option<CompositeTracePrincipal>,
principal_kind: PrincipalKind,
) -> Self {
Self {
scope,
allow_explain,
caller_identity,
conversation_id,
turn_id,
web_session_id,
control_fleet,
admin_fleet,
composite_trace,
principal_kind,
}
}
#[cfg(any(test, feature = "test-util"))]
pub fn set_scope_for_test(&mut self, scope: QueryScope) {
self.scope = scope;
}
#[cfg(any(test, feature = "test-util"))]
#[must_use]
pub const fn fleet_admin_for_test() -> Self {
Self {
scope: QueryScope::Fleet,
allow_explain: true,
caller_identity: None,
conversation_id: None,
turn_id: None,
web_session_id: None,
control_fleet: None,
admin_fleet: None,
composite_trace: None,
principal_kind: PrincipalKind::Admin,
}
}
#[cfg(any(test, feature = "test-util"))]
pub fn for_conversation_for_test(
conversation_id: &str,
subject: &GrantSubject,
) -> Result<Self, PrincipalError> {
Self::for_conversation(conversation_id, subject)
}
}
pub const SEARCH_SCOPE_CAP: usize = 5_000;
const SEARCH_SCOPE_HASH_DOMAIN: &[u8] = b"polychrome.search.scope.v1";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SearchScope {
conversation_ids: Vec<String>,
hash: String,
}
impl SearchScope {
pub(crate) const fn minted(conversation_ids: Vec<String>, hash: String) -> Self {
Self {
conversation_ids,
hash,
}
}
#[must_use]
pub fn conversation_ids(&self) -> &[String] {
&self.conversation_ids
}
#[must_use]
pub const fn count(&self) -> usize {
self.conversation_ids.len()
}
#[must_use]
pub fn hash(&self) -> &str {
&self.hash
}
}
#[derive(Debug, thiserror::Error)]
pub enum SearchScopeError {
#[error("persona store unavailable")]
StoreUnavailable,
#[error("persona is not active")]
PersonaNotActive,
#[error("participation count exceeds the search cap")]
OverCap {
count: usize,
},
}
pub(crate) fn canonical_search_scope(conversation_ids: Vec<String>) -> (Vec<String>, String) {
let mut encoded: Vec<(Vec<u8>, String)> = conversation_ids
.into_iter()
.map(|id| {
let bytes = id.as_bytes();
let len = u32::try_from(bytes.len())
.expect("conversation id byte length exceeds u32 — not an id this system mints");
let mut record = Vec::with_capacity(4 + bytes.len());
record.extend_from_slice(&len.to_be_bytes());
record.extend_from_slice(bytes);
(record, id)
})
.collect();
encoded.sort_by(|(a, _), (b, _)| a.cmp(b));
encoded.dedup_by(|(a, _), (b, _)| a == b);
let total_len = SEARCH_SCOPE_HASH_DOMAIN.len()
+ 1
+ encoded
.iter()
.map(|(record, _)| record.len())
.sum::<usize>();
let mut buf = Vec::with_capacity(total_len);
buf.extend_from_slice(SEARCH_SCOPE_HASH_DOMAIN);
buf.push(0);
for (record, _) in &encoded {
buf.extend_from_slice(record);
}
let hash = blake3::hash(&buf).to_hex().to_string();
let canonical_ids = encoded.into_iter().map(|(_, id)| id).collect();
(canonical_ids, hash)
}