stasis-rs 0.9.3

Durable AI orchestration framework with runtime jobs, lineage, and memory integration
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

use std::sync::Arc;

use async_trait::async_trait;
use serde_json::json;

use crate::application::orchestration::runtime_job_payloads::MemoryAggregateJobPayload;
use crate::application::runtime::in_memory_runtime::{JobExecutionOutcome, JobHandler};
use crate::application::runtime::memory_operation_job_outcome_helpers::{
    operation_failure, operation_success, policy_violation_failure,
};
use crate::application::runtime::memory_job_request_helpers::memory_scope_from_fields;
use crate::domain::errors::Result;
use crate::domain::runtime::job::Job;
use crate::ports::outbound::memory::memory_models::MemoryAggregateRequest;
use crate::ports::outbound::memory::memory_operations::MemoryOperations;

pub struct MemoryAggregateJobHandler {
    operations: Arc<dyn MemoryOperations>,
}

impl MemoryAggregateJobHandler {
    pub fn new(operations: Arc<dyn MemoryOperations>) -> Self {
        Self { operations }
    }

    fn parse_payload(raw: &str) -> std::result::Result<MemoryAggregateJobPayload, String> {
        serde_json::from_str(raw).map_err(|err| {
            format!("policy violation: invalid memory-aggregate payload json: {err}")
        })
    }
}

#[async_trait]
impl JobHandler for MemoryAggregateJobHandler {
    fn job_type(&self) -> &'static str {
        "workflow.stasis.memory.aggregate"
    }

    async fn execute(&self, job: &Job) -> Result<JobExecutionOutcome> {
        let payload = match Self::parse_payload(&job.payload_ref) {
            Ok(payload) => payload,
            Err(message) => return Ok(policy_violation_failure("stasis-memory-aggregate", message)),
        };

        let request = MemoryAggregateRequest {
            scope: memory_scope_from_fields(
                None,
                payload.session_ids,
                payload.tiers,
                payload.from_utc,
                payload.to_utc,
            ),
            max_groups: payload.max_groups.unwrap_or(30),
            max_nodes: payload.max_nodes.unwrap_or(5000),
        };

        match self.operations.aggregate(&request).await {
            Ok(result) => Ok(operation_success(
                "stasis-memory-aggregate",
                "memory-aggregate",
                &job.id,
                json!({
                    "total_groups": result.total_groups,
                    "scanned_nodes": result.scanned_nodes,
                }),
            )),
            Err(err) => Ok(operation_failure("stasis-memory-aggregate", err.to_string())),
        }
    }
}