Skip to main content

stasis/application/runtime/
memory_find_job_handler.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use serde_json::json;
5
6use crate::application::orchestration::runtime_job_payloads::MemoryFindJobPayload;
7use crate::application::runtime::in_memory_runtime::{JobExecutionOutcome, JobHandler};
8use crate::application::runtime::memory_operation_job_outcome_helpers::{
9    operation_failure, operation_success, policy_violation_failure,
10};
11use crate::domain::errors::Result;
12use crate::domain::runtime::job::Job;
13use crate::ports::outbound::memory::memory_context_reader::MemoryContextReader;
14use crate::ports::outbound::memory::memory_models::{
15    MemoryFilter, MemoryFindRequest, MemoryScope, MemorySortDirection, MemorySortField,
16};
17
18pub struct MemoryFindJobHandler {
19    reader: Arc<dyn MemoryContextReader>,
20}
21
22impl MemoryFindJobHandler {
23    pub fn new(reader: Arc<dyn MemoryContextReader>) -> Self {
24        Self { reader }
25    }
26
27    fn parse_payload(raw: &str) -> std::result::Result<MemoryFindJobPayload, String> {
28        serde_json::from_str(raw)
29            .map_err(|err| format!("policy violation: invalid memory-find payload json: {err}"))
30    }
31
32    fn map_sort_field(value: Option<&str>) -> MemorySortField {
33        match value {
34            Some("updated_at") => MemorySortField::UpdatedAt,
35            Some("psi") => MemorySortField::Psi,
36            Some("rho") => MemorySortField::Rho,
37            Some("kappa") => MemorySortField::Kappa,
38            _ => MemorySortField::Timestamp,
39        }
40    }
41
42    fn map_sort_direction(value: Option<&str>) -> MemorySortDirection {
43        match value {
44            Some("asc") => MemorySortDirection::Asc,
45            _ => MemorySortDirection::Desc,
46        }
47    }
48}
49
50#[async_trait]
51impl JobHandler for MemoryFindJobHandler {
52    fn job_type(&self) -> &'static str {
53        "workflow.stasis.memory.find"
54    }
55
56    async fn execute(&self, job: &Job) -> Result<JobExecutionOutcome> {
57        let payload = match Self::parse_payload(&job.payload_ref) {
58            Ok(payload) => payload,
59            Err(message) => return Ok(policy_violation_failure("stasis-memory-find", message)),
60        };
61
62        let request = MemoryFindRequest {
63            scope: MemoryScope {
64                session_ids: payload.session_ids,
65                tiers: payload.tiers,
66                from_utc: payload.from_utc,
67                to_utc: payload.to_utc,
68            },
69            filter: MemoryFilter {
70                text_contains: payload.text_contains,
71                ..Default::default()
72            },
73            limit: payload.limit.unwrap_or(50),
74            cursor: payload.cursor,
75            sort_field: Self::map_sort_field(payload.sort_field.as_deref()),
76            sort_direction: Self::map_sort_direction(payload.sort_direction.as_deref()),
77        };
78
79        match self.reader.find(&request).await {
80            Ok(result) => Ok(operation_success(
81                "stasis-memory-find",
82                "memory-find",
83                &job.id,
84                json!({
85                    "retrieved": result.retrieved,
86                    "has_more": result.has_more,
87                    "next_cursor": result.next_cursor,
88                    "node_sync_keys": result.node_sync_keys,
89                }),
90            )),
91            Err(err) => Ok(operation_failure("stasis-memory-find", err.to_string())),
92        }
93    }
94}