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