use std::collections::{BTreeSet, HashMap};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use meerkat::{
AttentionBindingRequest, AttentionListRequest, GoalAttentionTarget, WorkAttentionBinding,
WorkAttentionBindingId, WorkAttentionStatus, WorkAttentionTarget, WorkGraphError,
WorkGraphService, WorkNamespace, WorkOwnerKey, WorkOwnerKind,
};
pub const WORKGRAPH_ADMISSION_SIDECAR_FILE: &str = "workgraph.admission.sqlite3";
#[must_use]
pub fn workgraph_admission_sidecar_path(state_dir: &Path) -> PathBuf {
state_dir.join(WORKGRAPH_ADMISSION_SIDECAR_FILE)
}
pub type WorkGraphAdmissionSlot = Arc<std::sync::RwLock<Option<Arc<WorkGraphAdmission>>>>;
#[derive(Debug)]
pub(crate) enum WorkGraphAdmissionError {
Occupied { detail: String },
Service(WorkGraphError),
Lock(String),
}
pub(crate) struct WorkGraphAdmissionPermit {
_in_process: tokio::sync::OwnedMutexGuard<()>,
_cross_process: Option<SidecarLock>,
}
struct SidecarLock {
_connection: rusqlite::Connection,
}
impl SidecarLock {
const BUSY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
fn acquire(path: &Path) -> Result<Self, String> {
if !path.is_file() {
meerkat_sqlite::open(path, meerkat_sqlite::ConnectionProfile::PRIMARY)
.map_err(|error| format!("create admission sidecar {}: {error}", path.display()))?;
}
let connection = meerkat_sqlite::open_with(
path,
meerkat_sqlite::ConnectionProfile::Maintenance { write: true },
meerkat_sqlite::OpenOptions {
busy_timeout: Some(Self::BUSY_TIMEOUT),
..Default::default()
},
)
.map_err(|error| format!("open admission sidecar {}: {error}", path.display()))?;
connection
.execute_batch("BEGIN IMMEDIATE")
.map_err(|error| {
format!(
"could not lock the workgraph admission sidecar {} within the {}s busy \
timeout: {error}. The lock is held by another process sharing this state \
dir (in the documented deployment: a gateway and a library-mode runtime on \
one workgraph.sqlite3) — most likely a co-process is wedged mid-admission \
or under heavy binding-mutation load; retry, or check that co-process",
path.display(),
Self::BUSY_TIMEOUT.as_secs(),
)
})?;
Ok(Self {
_connection: connection,
})
}
}
pub struct WorkGraphAdmission {
mob_handle: meerkat_mob::MobHandle,
session_service: Option<Arc<dyn meerkat_mob::MobSessionService>>,
gate: Arc<tokio::sync::Mutex<()>>,
sidecar: Option<PathBuf>,
member_resolution_cache: std::sync::Mutex<
HashMap<meerkat::SessionId, (std::time::Instant, meerkat_mob::ids::AgentIdentity)>,
>,
member_resolution_ttl: std::time::Duration,
}
impl WorkGraphAdmission {
const MEMBER_RESOLUTION_CACHE_MAX: usize = 4096;
const MEMBER_RESOLUTION_TTL: std::time::Duration = std::time::Duration::from_mins(1);
pub fn new(
mob_handle: meerkat_mob::MobHandle,
session_service: Option<Arc<dyn meerkat_mob::MobSessionService>>,
sidecar: Option<PathBuf>,
) -> Self {
Self {
mob_handle,
session_service,
gate: Arc::new(tokio::sync::Mutex::new(())),
sidecar,
member_resolution_cache: std::sync::Mutex::new(HashMap::new()),
member_resolution_ttl: Self::MEMBER_RESOLUTION_TTL,
}
}
#[cfg(test)]
pub(crate) fn with_member_resolution_ttl(mut self, ttl: std::time::Duration) -> Self {
self.member_resolution_ttl = ttl;
self
}
pub(crate) fn mob_handle(&self) -> &meerkat_mob::MobHandle {
&self.mob_handle
}
async fn resolve_member_identity(
&self,
session_id: &meerkat::SessionId,
) -> Result<Option<meerkat_mob::ids::AgentIdentity>, WorkGraphError> {
if let Some(entry) = self
.mob_handle
.roster()
.await
.find_by_bridge_session_id(session_id)
{
return Ok(Some(entry.agent_identity.clone()));
}
if let Some((stamped_at, identity)) = self
.member_resolution_cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(session_id)
&& stamped_at.elapsed() < self.member_resolution_ttl
{
return Ok(Some(identity.clone()));
}
let Some(service) = self.session_service.as_ref() else {
return Ok(None);
};
let session = service
.load_persisted_session(session_id)
.await
.map_err(|error| {
WorkGraphError::Store(format!(
"workgraph admission could not read session {session_id} from the session \
store while resolving its mob member: {error}"
))
})?;
let resolved = session
.and_then(|session| session.session_metadata())
.and_then(|metadata| metadata.mob_member_binding)
.filter(|binding| binding.mob_id == self.mob_handle.definition().id.as_str())
.map(|binding| meerkat_mob::ids::AgentIdentity::from(binding.member.as_str()));
if let Some(identity) = resolved.as_ref() {
let mut cache = self
.member_resolution_cache
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if cache.len() >= Self::MEMBER_RESOLUTION_CACHE_MAX {
cache.clear();
}
cache.insert(
session_id.clone(),
(std::time::Instant::now(), identity.clone()),
);
}
Ok(resolved)
}
pub(crate) async fn lower_member_session_target(
&self,
target: GoalAttentionTarget,
) -> Result<GoalAttentionTarget, WorkGraphAdmissionError> {
let session_id = match &target {
GoalAttentionTarget::Session { session_id } => session_id.clone(),
GoalAttentionTarget::Owner { owner_key }
if owner_key.kind == WorkOwnerKind::Session =>
{
meerkat::SessionId::parse(&owner_key.id).map_err(|error| {
WorkGraphAdmissionError::Service(WorkGraphError::InvalidInput(format!(
"attention target owner key '{}' has kind 'session' but its id does not \
parse as a session id: {error}",
owner_key.canonical(),
)))
})?
}
_ => return Ok(target),
};
let Some(identity) = self
.resolve_member_identity(&session_id)
.await
.map_err(WorkGraphAdmissionError::Service)?
else {
return Ok(GoalAttentionTarget::Session { session_id });
};
Ok(
match meerkat_mob::lower_agent_identity_attention_target(
&self.mob_handle.definition().id,
&identity,
) {
Ok(lowered) => lowered,
Err(_) => GoalAttentionTarget::Session { session_id },
},
)
}
pub(crate) async fn acquire(
&self,
) -> Result<WorkGraphAdmissionPermit, WorkGraphAdmissionError> {
let in_process = Arc::clone(&self.gate).lock_owned().await;
let cross_process = match &self.sidecar {
None => None,
Some(path) => {
let path = path.clone();
let lock = tokio::task::spawn_blocking(move || SidecarLock::acquire(&path))
.await
.map_err(|error| {
WorkGraphAdmissionError::Lock(format!(
"admission sidecar lock task failed: {error}"
))
})?
.map_err(WorkGraphAdmissionError::Lock)?;
Some(lock)
}
};
Ok(WorkGraphAdmissionPermit {
_in_process: in_process,
_cross_process: cross_process,
})
}
pub(crate) async fn check_target_free(
&self,
service: &WorkGraphService,
namespace: Option<WorkNamespace>,
target: &WorkAttentionTarget,
exclude: Option<&WorkAttentionBindingId>,
action: &str,
) -> Result<(), WorkGraphAdmissionError> {
let aliases = self
.attention_target_alias_keys(target)
.await
.map_err(WorkGraphAdmissionError::Service)?;
let bindings = list_occupying_attention(service, namespace)
.await
.map_err(WorkGraphAdmissionError::Service)?;
let Some(existing) = bindings.iter().find(|binding| {
exclude != Some(&binding.binding_id)
&& binding_occupies_target(&binding.status)
&& binding
.target
.owner_key()
.is_ok_and(|key| aliases.contains(&key.canonical()))
}) else {
return Ok(());
};
let target_key = target
.owner_key()
.map_err(WorkGraphAdmissionError::Service)?;
Err(WorkGraphAdmissionError::Occupied {
detail: match existing.status {
WorkAttentionStatus::Paused { .. } => format!(
"target '{}' already has a paused attention binding {} that will reactivate \
when its pause expires; resume it or close its goal instead of {action}",
target_key.canonical(),
existing.binding_id,
),
_ => format!(
"target '{}' already has an active attention binding {}; reassign it or \
close its goal before {action}",
target_key.canonical(),
existing.binding_id,
),
},
})
}
pub(crate) async fn check_resume_target_free(
&self,
service: &WorkGraphService,
namespace: Option<WorkNamespace>,
binding_id: &WorkAttentionBindingId,
) -> Result<(), WorkGraphAdmissionError> {
let resumed = match service
.attention_binding(AttentionBindingRequest {
binding_id: binding_id.clone(),
realm_id: None,
namespace: namespace.clone(),
})
.await
{
Ok(result) => result.attention,
Err(WorkGraphError::AttentionNotFound { .. }) => return Ok(()),
Err(error) => return Err(WorkGraphAdmissionError::Service(error)),
};
let aliases = self
.attention_target_alias_keys(&resumed.target)
.await
.map_err(WorkGraphAdmissionError::Service)?;
let siblings = list_occupying_attention(service, namespace)
.await
.map_err(WorkGraphAdmissionError::Service)?;
let Some(other) = siblings.iter().find(|binding| {
binding.binding_id != *binding_id
&& binding_occupies_target(&binding.status)
&& binding
.target
.owner_key()
.is_ok_and(|key| aliases.contains(&key.canonical()))
}) else {
return Ok(());
};
let target_key = resumed
.target
.owner_key()
.map(|key| key.canonical())
.unwrap_or_default();
Err(WorkGraphAdmissionError::Occupied {
detail: match other.status {
WorkAttentionStatus::Paused { .. } => format!(
"resuming attention binding {binding_id} would give target '{target_key}' a \
second occupying binding: {} is paused and will reactivate when its pause \
expires; close its goal first",
other.binding_id,
),
_ => format!(
"resuming attention binding {binding_id} would give target '{target_key}' a \
second active binding ({} is already active); reassign it or close its \
goal first",
other.binding_id,
),
},
})
}
async fn attention_target_alias_keys(
&self,
target: &WorkAttentionTarget,
) -> Result<BTreeSet<String>, WorkGraphError> {
let mob_handle = &self.mob_handle;
let primary = target.owner_key()?;
let mut keys = BTreeSet::from([primary.canonical()]);
match primary.kind {
WorkOwnerKind::Session => {
if let Ok(session_id) = meerkat::SessionId::parse(&primary.id)
&& let Some(identity) = self.resolve_member_identity(&session_id).await?
&& let Ok(key) = meerkat_mob::lower_agent_identity_owner_key(
&mob_handle.definition().id,
&identity,
)
{
keys.insert(key.canonical());
}
}
WorkOwnerKind::Agent => {
if let Some((mob_id, identity)) = mob_agent_owner_key_parts(&primary.id)
&& mob_id == mob_handle.definition().id.as_str()
&& let Some(session_id) = mob_handle
.resolve_bridge_session_id_observation(
&meerkat_mob::ids::AgentIdentity::from(identity),
)
.await
&& let Ok(key) = WorkOwnerKey::session(session_id.to_string())
{
keys.insert(key.canonical());
}
}
_ => {}
}
Ok(keys)
}
}
fn binding_occupies_target(status: &WorkAttentionStatus) -> bool {
matches!(
status,
WorkAttentionStatus::Active | WorkAttentionStatus::Paused { .. }
)
}
async fn list_occupying_attention(
service: &WorkGraphService,
namespace: Option<WorkNamespace>,
) -> Result<Vec<WorkAttentionBinding>, WorkGraphError> {
let namespace = namespace.unwrap_or_else(|| service.default_namespace().clone());
let mut bindings = Vec::new();
for status in [
WorkAttentionStatus::Active,
WorkAttentionStatus::Paused { until: None },
] {
let result = service
.list_attention(AttentionListRequest {
realm_id: Some(service.default_realm_id().to_string()),
namespace: Some(namespace.clone()),
target: None,
status: Some(status),
})
.await?;
bindings.extend(result.attention);
}
Ok(bindings)
}
fn mob_agent_owner_key_parts(owner_id: &str) -> Option<(&str, &str)> {
let rest = owner_id.strip_prefix("mob/")?;
let (mob_id, agent_identity) = rest.split_once("/agent/")?;
if mob_id.is_empty()
|| agent_identity.is_empty()
|| mob_id.contains('/')
|| agent_identity.contains('/')
{
return None;
}
Some((mob_id, agent_identity))
}
#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)]
mod tests {
use super::*;
#[tokio::test(flavor = "multi_thread")]
async fn sidecar_lock_admits_one_holder_and_makes_the_second_wait() {
let dir = tempfile::tempdir().expect("temp dir");
let path = workgraph_admission_sidecar_path(dir.path());
let first = SidecarLock::acquire(&path).expect("first lock");
assert!(path.exists(), "acquire must create the sidecar database");
let contended = path.clone();
let second = tokio::task::spawn_blocking(move || SidecarLock::acquire(&contended));
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
assert!(
!second.is_finished(),
"second holder must wait while the first transaction is open"
);
drop(first);
let second = second.await.expect("join");
assert!(second.is_ok(), "released lock must admit the waiter");
}
#[test]
fn sidecar_carries_no_schema_and_no_ledger() {
let dir = tempfile::tempdir().expect("temp dir");
let path = workgraph_admission_sidecar_path(dir.path());
drop(SidecarLock::acquire(&path).expect("acquire"));
let probe = rusqlite::Connection::open(&path).expect("probe");
let tables: i64 = probe
.query_row(
"SELECT COUNT(*) FROM sqlite_master WHERE type = 'table'",
[],
|row| row.get(0),
)
.expect("count tables");
assert_eq!(
tables, 0,
"the lock database must stay empty: no ledger, no tables"
);
}
#[test]
fn sidecar_is_a_separate_file_from_the_store() {
assert_eq!(
WORKGRAPH_ADMISSION_SIDECAR_FILE,
"workgraph.admission.sqlite3"
);
assert_ne!(
WORKGRAPH_ADMISSION_SIDECAR_FILE,
crate::workgraph_wiring::WORKGRAPH_STORE_FILE
);
let dir = Path::new("/state");
assert_eq!(
workgraph_admission_sidecar_path(dir),
dir.join("workgraph.admission.sqlite3")
);
}
}