systemprompt_provider_contracts/

job.rs

//! [`Job`] contract for scheduled / on-startup background jobs registered
//! via the `inventory` crate.
//!
//! Jobs are dispatched as `&'static dyn Job` from the inventory, so the trait
//! uses `#[async_trait]`; native `async fn` in traits is not `dyn`-compatible.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use std::collections::HashMap;
use std::sync::Arc;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use systemprompt_identifiers::Actor;

use crate::error::ProviderResult;

#[derive(Debug, Clone)]
pub struct JobResult {
    pub success: bool,
    pub message: Option<String>,
    pub items_processed: Option<u64>,
    pub items_failed: Option<u64>,
    pub duration_ms: u64,
}

impl JobResult {
    #[must_use]
    pub const fn success() -> Self {
        Self {
            success: true,
            message: None,
            items_processed: None,
            items_failed: None,
            duration_ms: 0,
        }
    }

    #[must_use]
    pub fn with_message(mut self, message: impl Into<String>) -> Self {
        self.message = Some(message.into());
        self
    }

    #[must_use]
    pub const fn with_stats(mut self, processed: u64, failed: u64) -> Self {
        self.items_processed = Some(processed);
        self.items_failed = Some(failed);
        self
    }

    #[must_use]
    pub const fn with_duration(mut self, duration_ms: u64) -> Self {
        self.duration_ms = duration_ms;
        self
    }

    #[must_use]
    pub fn failure(message: impl Into<String>) -> Self {
        Self {
            success: false,
            message: Some(message.into()),
            items_processed: None,
            items_failed: None,
            duration_ms: 0,
        }
    }
}

pub struct JobContext {
    actor: Actor,
    db_pool: Arc<dyn std::any::Any + Send + Sync>,
    app_context: Arc<dyn std::any::Any + Send + Sync>,
    app_paths: Arc<dyn std::any::Any + Send + Sync>,
    parameters: HashMap<String, String>,
    enforce: bool,
}

impl std::fmt::Debug for JobContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("JobContext")
            .field("actor", &self.actor)
            .field("db_pool", &"<type-erased>")
            .field("app_context", &"<type-erased>")
            .field("app_paths", &"<type-erased>")
            .field("parameters", &self.parameters)
            .field("enforce", &self.enforce)
            .finish()
    }
}

impl JobContext {
    #[must_use]
    pub fn new(
        actor: Actor,
        db_pool: Arc<dyn std::any::Any + Send + Sync>,
        app_context: Arc<dyn std::any::Any + Send + Sync>,
        app_paths: Arc<dyn std::any::Any + Send + Sync>,
    ) -> Self {
        Self {
            actor,
            db_pool,
            app_context,
            app_paths,
            parameters: HashMap::new(),
            enforce: false,
        }
    }

    #[must_use]
    pub const fn enforce(&self) -> bool {
        self.enforce
    }

    #[must_use]
    pub const fn with_enforce(mut self, enforce: bool) -> Self {
        self.enforce = enforce;
        self
    }

    #[must_use]
    pub const fn actor(&self) -> &Actor {
        &self.actor
    }

    #[must_use]
    pub fn with_parameters(mut self, parameters: HashMap<String, String>) -> Self {
        self.parameters = parameters;
        self
    }

    #[must_use]
    pub fn db_pool<T: 'static>(&self) -> Option<&T> {
        self.db_pool.as_ref().downcast_ref::<T>()
    }

    #[must_use]
    pub fn app_context<T: 'static>(&self) -> Option<&T> {
        self.app_context.as_ref().downcast_ref::<T>()
    }

    #[must_use]
    pub fn app_paths<T: 'static>(&self) -> Option<&T> {
        self.app_paths.as_ref().downcast_ref::<T>()
    }

    #[must_use]
    pub fn db_pool_arc(&self) -> Arc<dyn std::any::Any + Send + Sync> {
        Arc::clone(&self.db_pool)
    }

    #[must_use]
    pub fn app_context_arc(&self) -> Arc<dyn std::any::Any + Send + Sync> {
        Arc::clone(&self.app_context)
    }

    #[must_use]
    pub fn app_paths_arc(&self) -> Arc<dyn std::any::Any + Send + Sync> {
        Arc::clone(&self.app_paths)
    }

    #[must_use]
    pub const fn parameters(&self) -> &HashMap<String, String> {
        &self.parameters
    }

    #[must_use]
    pub fn get_parameter(&self, key: &str) -> Option<&String> {
        self.parameters.get(key)
    }

    pub fn get_parameter_parsed<T: std::str::FromStr>(
        &self,
        key: &str,
    ) -> Result<Option<T>, crate::ProviderError>
    where
        T::Err: std::fmt::Display,
    {
        self.parameters
            .get(key)
            .map(|value| {
                value.parse().map_err(|e| {
                    crate::ProviderError::Configuration(format!(
                        "invalid job parameter {key}={value}: {e}"
                    ))
                })
            })
            .transpose()
    }
}

/// Where a job runs when the scheduler has several replicas.
///
/// `Cluster` jobs run once per tick across the whole deployment (one replica
/// wins the advisory lock). `Node` jobs run on every replica, because their
/// effect is local to the process or its filesystem.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum JobScope {
    #[default]
    Cluster,
    Node,
}

#[async_trait]
pub trait Job: Send + Sync + 'static {
    fn name(&self) -> &'static str;

    fn description(&self) -> &'static str {
        ""
    }

    fn schedule(&self) -> &'static str;

    fn tags(&self) -> Vec<&'static str> {
        vec![]
    }

    async fn execute(&self, ctx: &JobContext) -> ProviderResult<JobResult>;

    fn enabled(&self) -> bool {
        true
    }

    fn schedulable(&self) -> bool {
        true
    }

    fn scope(&self) -> JobScope {
        JobScope::Cluster
    }
}

inventory::collect!(&'static dyn Job);

#[macro_export]
macro_rules! submit_job {
    ($job:expr) => {
        inventory::submit!($job as &'static dyn $crate::Job);
    };
}