Skip to main content

stasis/application/runtime/
memory_recall_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::{
7    MemoryPolicyPayload, MemoryRecallJobPayload,
8};
9use crate::application::runtime::in_memory_runtime::{JobExecutionOutcome, JobHandler};
10use crate::application::runtime::memory_recall_request_builder::build_memory_recall_request;
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::MemoryRecallRequest;
15
16pub struct MemoryRecallJobHandler {
17    reader: Arc<dyn MemoryContextReader>,
18}
19
20impl MemoryRecallJobHandler {
21    pub fn new(reader: Arc<dyn MemoryContextReader>) -> Self {
22        Self { reader }
23    }
24
25    fn parse_payload(raw: &str) -> std::result::Result<MemoryRecallJobPayload, String> {
26        serde_json::from_str(raw)
27            .map_err(|err| format!("policy violation: invalid memory-recall payload json: {err}"))
28    }
29
30    fn build_request(
31        correlation_id: &str,
32        policy: Option<&MemoryPolicyPayload>,
33    ) -> MemoryRecallRequest {
34        build_memory_recall_request(correlation_id, None, policy)
35    }
36}
37
38#[async_trait]
39impl JobHandler for MemoryRecallJobHandler {
40    fn job_type(&self) -> &'static str {
41        "workflow.stasis.memory.recall"
42    }
43
44    async fn execute(&self, job: &Job) -> Result<JobExecutionOutcome> {
45        let payload = match Self::parse_payload(&job.payload_ref) {
46            Ok(payload) => payload,
47            Err(message) => {
48                let diagnostics = json!({
49                    "provider": "stasis-memory-recall",
50                    "status": "failure",
51                    "guardrail_code": "POLICY_VIOLATION",
52                    "policy_reason": message,
53                })
54                .to_string();
55                return Ok(JobExecutionOutcome::FatalFailure {
56                    message: "invalid memory recall payload".to_string(),
57                    execution_id: None,
58                    diagnostics: Some(diagnostics),
59                });
60            }
61        };
62
63        let recall_request =
64            Self::build_request(&job.correlation_id, payload.memory_policy.as_ref());
65        match self.reader.recall(&recall_request).await {
66            Ok(response) => Ok(JobExecutionOutcome::Success {
67                sttp_output_node_id: format!("sttp:memory-recall:{}", job.id),
68                execution_id: None,
69                diagnostics: Some(
70                    json!({
71                        "provider": "stasis-memory-recall",
72                        "status": "success",
73                        "retrieved": response.retrieved,
74                        "retrieval_path": response.retrieval_path,
75                        "fallback_triggered": response.fallback_triggered,
76                        "fallback_reason": response.fallback_reason,
77                        "has_more": response.has_more,
78                    })
79                    .to_string(),
80                ),
81            }),
82            Err(err) => Ok(JobExecutionOutcome::FatalFailure {
83                message: err.to_string(),
84                execution_id: None,
85                diagnostics: Some(
86                    json!({
87                        "provider": "stasis-memory-recall",
88                        "status": "failure",
89                        "error": err.to_string(),
90                    })
91                    .to_string(),
92                ),
93            }),
94        }
95    }
96}