use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use chrono::Duration as ChronoDuration;
use meerkat_core::{ContentInput, SessionId, skills::SkillRef, types::RenderMetadata};
use meerkat_runtime::CompletionHandle;
use meerkat_schedule::{
DeliveryCompletion, DeliveryDispatch, DeliveryReceipt, DeliveryReceiptStage, DeliveryTerminal,
MobTargetBinding, Occurrence, OccurrenceFailureClass, OccurrencePhase, ScheduleDomainError,
ScheduleDriver, ScheduleDriverConfig, ScheduleService, ScheduleStoreKind,
ScheduleTargetDelivery, ScheduleTargetProbe, ScheduledSessionAction,
SessionMaterializationSpec, SessionTargetBinding, TargetBinding, TargetProbeOutcome,
};
#[cfg(not(target_arch = "wasm32"))]
use tokio::sync::oneshot;
#[cfg(not(target_arch = "wasm32"))]
use tokio::task::JoinHandle;
#[cfg(target_arch = "wasm32")]
use tokio_with_wasm::alias::sync::oneshot;
#[cfg(target_arch = "wasm32")]
use tokio_with_wasm::alias::task::JoinHandle;
pub struct ScheduleHostHandle {
shutdown_tx: Option<oneshot::Sender<()>>,
join: JoinHandle<()>,
}
impl ScheduleHostHandle {
pub async fn shutdown(mut self) {
if let Some(shutdown_tx) = self.shutdown_tx.take() {
let _ = shutdown_tx.send(());
}
let _ = self.join.await;
}
}
#[derive(Debug, Clone)]
struct ResolvedScheduledSession {
session_id: SessionId,
materialized_session_id: Option<SessionId>,
allow_system_prompt_override: bool,
}
pub struct AcceptedScheduledInput {
pub correlation_id: Option<String>,
pub handle: Option<CompletionHandle>,
}
#[derive(Debug, Clone)]
pub struct ScheduledPromptDispatch {
pub prompt: ContentInput,
pub render_metadata: Option<RenderMetadata>,
pub skill_refs: Vec<SkillRef>,
pub additional_instructions: Vec<String>,
pub materialized_session_id: Option<SessionId>,
}
#[async_trait]
pub trait SurfaceScheduleSessionHost: Send + Sync {
async fn probe_session_target(
&self,
binding: &SessionTargetBinding,
) -> Result<TargetProbeOutcome, ScheduleDomainError>;
async fn materialize_session(
&self,
create: &SessionMaterializationSpec,
prompt_system_prompt: Option<&str>,
) -> Result<SessionId, ScheduleDomainError>;
async fn deliver_prompt(
&self,
session_id: &SessionId,
occurrence: &Occurrence,
dispatch: ScheduledPromptDispatch,
) -> Result<DeliveryDispatch, ScheduleDomainError>;
async fn deliver_event(
&self,
session_id: &SessionId,
occurrence: &Occurrence,
event_type: String,
payload: serde_json::Value,
render_metadata: Option<RenderMetadata>,
materialized_session_id: Option<SessionId>,
) -> Result<DeliveryDispatch, ScheduleDomainError>;
}
#[async_trait]
pub trait SurfaceScheduleMobHost: Send + Sync {
async fn probe_mob_target(
&self,
binding: &MobTargetBinding,
) -> Result<TargetProbeOutcome, ScheduleDomainError>;
async fn deliver_mob_target(
&self,
occurrence: &Occurrence,
binding: &MobTargetBinding,
) -> Result<DeliveryDispatch, ScheduleDomainError>;
}
pub struct NoopScheduleMobHost {
detail: String,
}
impl NoopScheduleMobHost {
pub fn new(detail: impl Into<String>) -> Self {
Self {
detail: detail.into(),
}
}
}
#[async_trait]
impl SurfaceScheduleMobHost for NoopScheduleMobHost {
async fn probe_mob_target(
&self,
_binding: &MobTargetBinding,
) -> Result<TargetProbeOutcome, ScheduleDomainError> {
Ok(TargetProbeOutcome::Missing {
detail: Some(self.detail.clone()),
})
}
async fn deliver_mob_target(
&self,
occurrence: &Occurrence,
_binding: &MobTargetBinding,
) -> Result<DeliveryDispatch, ScheduleDomainError> {
Ok(immediate_delivery_failure(
occurrence,
self.detail.clone(),
OccurrenceFailureClass::MobRejected,
None,
None,
))
}
}
pub struct SharedScheduleTargetAdapter {
schedule_service: ScheduleService,
session_host: Arc<dyn SurfaceScheduleSessionHost>,
mob_host: Arc<dyn SurfaceScheduleMobHost>,
}
impl SharedScheduleTargetAdapter {
pub fn new(
schedule_service: ScheduleService,
session_host: Arc<dyn SurfaceScheduleSessionHost>,
mob_host: Arc<dyn SurfaceScheduleMobHost>,
) -> Self {
Self {
schedule_service,
session_host,
mob_host,
}
}
async fn resolve_session(
&self,
occurrence: &Occurrence,
binding: &SessionTargetBinding,
) -> Result<ResolvedScheduledSession, DeliveryDispatch> {
match binding {
SessionTargetBinding::ExactSession { session_id, .. }
| SessionTargetBinding::ResumableSession { session_id, .. } => {
Ok(ResolvedScheduledSession {
session_id: session_id.clone(),
materialized_session_id: None,
allow_system_prompt_override: false,
})
}
SessionTargetBinding::MaterializeOnDemandSession {
bound_session_id: Some(session_id),
..
} => Ok(ResolvedScheduledSession {
session_id: session_id.clone(),
materialized_session_id: Some(session_id.clone()),
allow_system_prompt_override: false,
}),
SessionTargetBinding::MaterializeOnDemandSession {
create,
action,
bound_session_id: None,
} => {
let prompt_system_prompt = match action {
ScheduledSessionAction::Prompt { system_prompt, .. } => {
system_prompt.as_deref()
}
ScheduledSessionAction::Event { .. } => None,
};
match self
.session_host
.materialize_session(create, prompt_system_prompt)
.await
{
Ok(session_id) => {
if let Err(error) = self
.schedule_service
.bind_materialized_session_for_occurrence(occurrence, &session_id)
.await
{
return Err(immediate_delivery_failure(
occurrence,
error.to_string(),
OccurrenceFailureClass::InternalError,
None,
Some(session_id),
));
}
Ok(ResolvedScheduledSession {
session_id: session_id.clone(),
materialized_session_id: Some(session_id),
allow_system_prompt_override: true,
})
}
Err(error) => Err(immediate_delivery_failure(
occurrence,
error.to_string(),
OccurrenceFailureClass::TargetMaterializationFailed,
None,
None,
)),
}
}
}
}
}
#[async_trait]
impl ScheduleTargetProbe for SharedScheduleTargetAdapter {
async fn probe_target(
&self,
occurrence: &Occurrence,
) -> Result<TargetProbeOutcome, ScheduleDomainError> {
match &occurrence.target_snapshot {
TargetBinding::Session(binding) => {
self.session_host.probe_session_target(binding).await
}
TargetBinding::Mob(binding) => self.mob_host.probe_mob_target(binding).await,
}
}
}
#[async_trait]
impl ScheduleTargetDelivery for SharedScheduleTargetAdapter {
async fn deliver_occurrence(
&self,
occurrence: &Occurrence,
) -> Result<DeliveryDispatch, ScheduleDomainError> {
match &occurrence.target_snapshot {
TargetBinding::Session(binding) => {
let resolved = match self.resolve_session(occurrence, binding).await {
Ok(resolved) => resolved,
Err(dispatch) => return Ok(dispatch),
};
match binding.action() {
ScheduledSessionAction::Prompt {
prompt,
system_prompt,
render_metadata,
skill_refs,
additional_instructions,
} => {
if system_prompt.is_some() && !resolved.allow_system_prompt_override {
return Ok(immediate_delivery_failure(
occurrence,
"scheduled system_prompt override is only supported when materializing a new session"
.to_string(),
OccurrenceFailureClass::RuntimeRejected,
None,
resolved.materialized_session_id,
));
}
self.session_host
.deliver_prompt(
&resolved.session_id,
occurrence,
ScheduledPromptDispatch {
prompt: prompt.clone(),
render_metadata: render_metadata.clone(),
skill_refs: skill_refs.clone(),
additional_instructions: additional_instructions.clone(),
materialized_session_id: resolved.materialized_session_id,
},
)
.await
}
ScheduledSessionAction::Event {
event_type,
payload,
render_metadata,
} => {
self.session_host
.deliver_event(
&resolved.session_id,
occurrence,
event_type.clone(),
payload.clone(),
render_metadata.clone(),
resolved.materialized_session_id,
)
.await
}
}
}
TargetBinding::Mob(binding) => {
self.mob_host.deliver_mob_target(occurrence, binding).await
}
}
}
}
pub fn schedule_host_supported(kind: ScheduleStoreKind) -> bool {
!matches!(kind, ScheduleStoreKind::Disabled | ScheduleStoreKind::Jsonl)
}
pub fn spawn_schedule_host(
schedule_service: ScheduleService,
adapter: Arc<SharedScheduleTargetAdapter>,
owner_id: impl Into<String>,
) -> ScheduleHostHandle {
let driver = Arc::new(ScheduleDriver::new(
schedule_service.clone(),
schedule_service.store(),
adapter.clone(),
adapter,
owner_id,
ScheduleDriverConfig {
claim_limit: 32,
lease_duration: ChronoDuration::seconds(60),
},
));
let poll_interval = if cfg!(test) {
Duration::from_millis(50)
} else {
Duration::from_millis(250)
};
let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
#[cfg(not(target_arch = "wasm32"))]
let join = tokio::spawn(async move {
let mut interval = tokio::time::interval(poll_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tokio::select! {
_ = &mut shutdown_rx => break,
_ = interval.tick() => {
let _ = driver.tick_once().await;
}
}
}
});
#[cfg(target_arch = "wasm32")]
let join = tokio_with_wasm::alias::task::spawn(async move {
let mut interval = tokio_with_wasm::alias::time::interval(poll_interval);
loop {
tokio_with_wasm::alias::select! {
_ = &mut shutdown_rx => break,
() = interval.tick() => {
let _ = driver.tick_once().await;
}
}
}
});
ScheduleHostHandle {
shutdown_tx: Some(shutdown_tx),
join,
}
}
pub fn build_dispatch_from_accepted(
occurrence: &Occurrence,
accepted: AcceptedScheduledInput,
materialized_session_id: Option<SessionId>,
) -> DeliveryDispatch {
let mut receipt = DeliveryReceipt::new(
occurrence.occurrence_id.clone(),
occurrence.attempt_count,
DeliveryReceiptStage::DispatchAccepted,
);
receipt.correlation_id = accepted.correlation_id.clone();
receipt.materialized_session_id = materialized_session_id.clone();
let occurrence_id = occurrence.occurrence_id.clone();
let attempt_count = occurrence.attempt_count;
let correlation_id = accepted.correlation_id.clone();
let completed_materialized_session_id = materialized_session_id.clone();
let _ = accepted.handle;
let completion = Box::pin(async move {
let mut receipt = DeliveryReceipt::new(
occurrence_id,
attempt_count,
DeliveryReceiptStage::Completed,
);
receipt.correlation_id = correlation_id;
receipt.materialized_session_id = completed_materialized_session_id;
Ok(DeliveryTerminal::completed(Some(receipt)))
});
DeliveryDispatch {
receipt,
correlation_id: accepted.correlation_id,
materialized_session_id,
completion,
}
}
pub fn immediate_completed_dispatch(
occurrence: &Occurrence,
correlation_id: Option<String>,
) -> DeliveryDispatch {
let mut receipt = DeliveryReceipt::new(
occurrence.occurrence_id.clone(),
occurrence.attempt_count,
DeliveryReceiptStage::DispatchAccepted,
);
receipt.correlation_id = correlation_id.clone();
DeliveryDispatch {
receipt,
correlation_id,
materialized_session_id: None,
completion: Box::pin(async { Ok(DeliveryTerminal::completed(None)) }),
}
}
pub fn async_completion_dispatch(
occurrence: &Occurrence,
correlation_id: Option<String>,
completion: DeliveryCompletion,
) -> DeliveryDispatch {
let mut receipt = DeliveryReceipt::new(
occurrence.occurrence_id.clone(),
occurrence.attempt_count,
DeliveryReceiptStage::DispatchAccepted,
);
receipt.correlation_id = correlation_id.clone();
DeliveryDispatch {
receipt,
correlation_id,
materialized_session_id: None,
completion,
}
}
pub fn immediate_delivery_failure(
occurrence: &Occurrence,
detail: String,
failure_class: OccurrenceFailureClass,
correlation_id: Option<String>,
materialized_session_id: Option<SessionId>,
) -> DeliveryDispatch {
let mut receipt = DeliveryReceipt::new(
occurrence.occurrence_id.clone(),
occurrence.attempt_count,
DeliveryReceiptStage::DispatchStarted,
);
receipt.correlation_id = correlation_id.clone();
receipt.materialized_session_id = materialized_session_id.clone();
DeliveryDispatch {
receipt,
correlation_id,
materialized_session_id: materialized_session_id.clone(),
completion: Box::pin(async move {
Ok(DeliveryTerminal {
phase: OccurrencePhase::DeliveryFailed,
receipt: None,
detail: Some(detail),
failure_class: Some(failure_class),
runtime_outcome: None,
materialized_session_id,
})
}),
}
}
pub fn schedule_attempt_idempotency_key(occurrence: &Occurrence) -> String {
format!(
"schedule:{}:occurrence:{}:attempt:{}",
occurrence.schedule_id, occurrence.occurrence_id, occurrence.attempt_count
)
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use std::collections::BTreeMap;
fn sample_occurrence() -> Occurrence {
let schedule = meerkat_schedule::Schedule::new(meerkat_schedule::CreateScheduleRequest {
name: Some("schedule-host-test".to_string()),
description: None,
trigger: meerkat_schedule::TriggerSpec::Interval(
meerkat_schedule::IntervalTriggerSpec {
start_at_utc: chrono::Utc::now(),
every_seconds: 60,
end_at_utc: None,
},
),
target: TargetBinding::session(SessionTargetBinding::ExactSession {
session_id: SessionId::new(),
action: ScheduledSessionAction::Prompt {
prompt: ContentInput::Text("hello".to_string()),
system_prompt: None,
render_metadata: None,
skill_refs: Vec::new(),
additional_instructions: Vec::new(),
},
}),
misfire_policy: meerkat_schedule::MisfirePolicy::Skip,
overlap_policy: meerkat_schedule::OverlapPolicy::SkipIfRunning,
missing_target_policy: meerkat_schedule::MissingTargetPolicy::Skip,
labels: BTreeMap::new(),
planning_horizon_days: None,
planning_horizon_occurrences: None,
});
let mut occurrence = Occurrence::planned_from_schedule(
&schedule,
meerkat_schedule::OccurrenceOrdinal(0),
chrono::Utc::now(),
);
occurrence.attempt_count = 1;
occurrence
}
#[tokio::test]
async fn noop_mob_host_reports_clear_feature_required_failure() {
let host = NoopScheduleMobHost::new(
"scheduled mob targets require the mob feature on the CLI host",
);
let binding = MobTargetBinding::Member {
mob_id: "ops".to_string(),
member_id: "deploy-monitor".to_string(),
action: meerkat_schedule::ScheduledMobAction::Send {
content: ContentInput::Text("Check deploy state.".to_string()),
render_metadata: None,
},
};
let probe = host
.probe_mob_target(&binding)
.await
.expect("probe should succeed");
let TargetProbeOutcome::Missing { detail } = probe else {
panic!("expected no-op mob host to report missing, got {probe:?}");
};
assert_eq!(
detail.as_deref(),
Some("scheduled mob targets require the mob feature on the CLI host")
);
let occurrence = sample_occurrence();
let dispatch = host
.deliver_mob_target(&occurrence, &binding)
.await
.expect("delivery dispatch");
let terminal = dispatch.completion.await.expect("delivery terminal");
assert_eq!(terminal.phase, OccurrencePhase::DeliveryFailed);
assert_eq!(
terminal.detail.as_deref(),
Some("scheduled mob targets require the mob feature on the CLI host")
);
assert_eq!(
terminal.failure_class,
Some(OccurrenceFailureClass::MobRejected)
);
}
}