stasis/application/runtime/
prompt_chat_job_handler.rs1use std::sync::Arc;
2
3use async_trait::async_trait;
4use serde_json::json;
5
6use crate::application::orchestration::runtime_job_payloads::{
7 PromptJobPayload,
8};
9use crate::application::runtime::identity_context_compiler::{
10 load_identity_context_summary, prepend_identity_snapshot,
11};
12use crate::application::runtime::memory_recall_context_compiler::prepend_memory_recall_context;
13use crate::application::runtime::memory_persistence_helpers::{
14 SttpPromptNodeFormat, memory_query_fingerprint, memory_query_id, memory_scope_hash,
15 render_prompt_response_sttp_node, resolve_sttp_output_node_id, should_store,
16};
17use crate::application::runtime::memory_recall_request_builder::build_memory_recall_request;
18use crate::application::orchestration::prompt_pipeline::{
19 PromptExecutionPipeline, PromptExecutionRequest,
20};
21use crate::application::runtime::in_memory_runtime::{JobExecutionOutcome, JobHandler};
22use crate::application::runtime::runtime_diagnostics_helpers::{
23 build_runtime_failure_identity_context_section, build_runtime_failure_memory_recall_section,
24 build_runtime_memory_diagnostics_bundle, RuntimeIdentityDiagnosticsInput,
25 RuntimeMemoryRecallDiagnosticsInput, RuntimeMemoryStoreDiagnosticsInput,
26};
27use crate::application::runtime::runtime_handler_execution_context::RuntimeHandlerExecutionContext;
28use crate::domain::errors::Result;
29use crate::domain::runtime::job::Job;
30use crate::ports::outbound::ai_chat_client::AiChatClient;
31use crate::ports::outbound::memory::identity_memory_store::IdentityMemoryStore;
32use crate::ports::outbound::memory::memory_context_reader::MemoryContextReader;
33use crate::ports::outbound::memory::memory_context_writer::MemoryContextWriter;
34use crate::ports::outbound::memory::memory_models::MemoryStoreRequest;
35
36pub struct PromptChatJobHandler {
37 pipeline: PromptExecutionPipeline,
38 memory_reader: Option<Arc<dyn MemoryContextReader>>,
39 memory_writer: Option<Arc<dyn MemoryContextWriter>>,
40 identity_memory_store: Option<Arc<dyn IdentityMemoryStore>>,
41}
42
43impl PromptChatJobHandler {
44 pub fn new(chat_client: Arc<dyn AiChatClient>) -> Self {
45 Self::new_with_memory_and_identity(chat_client, None, None, None)
46 }
47
48 pub fn new_with_memory(
49 chat_client: Arc<dyn AiChatClient>,
50 memory_reader: Option<Arc<dyn MemoryContextReader>>,
51 memory_writer: Option<Arc<dyn MemoryContextWriter>>,
52 ) -> Self {
53 Self::new_with_memory_and_identity(chat_client, memory_reader, memory_writer, None)
54 }
55
56 pub fn new_with_memory_and_identity(
57 chat_client: Arc<dyn AiChatClient>,
58 memory_reader: Option<Arc<dyn MemoryContextReader>>,
59 memory_writer: Option<Arc<dyn MemoryContextWriter>>,
60 identity_memory_store: Option<Arc<dyn IdentityMemoryStore>>,
61 ) -> Self {
62 Self {
63 pipeline: PromptExecutionPipeline::new(chat_client),
64 memory_reader,
65 memory_writer,
66 identity_memory_store,
67 }
68 }
69
70 fn parse_payload(raw: &str) -> std::result::Result<PromptJobPayload, String> {
71 let payload: PromptJobPayload = serde_json::from_str(raw)
72 .map_err(|err| format!("policy violation: invalid prompt job payload json: {err}"))?;
73
74 if payload.user_prompt.trim().is_empty() {
75 return Err(
76 "policy violation: prompt payload.user_prompt must be non-empty".to_string(),
77 );
78 }
79
80 Ok(payload)
81 }
82
83 fn build_failure(message: String) -> JobExecutionOutcome {
84 let diagnostics = json!({
85 "provider": "stasis-pipeline",
86 "status": "failure",
87 "guardrail_code": "POLICY_VIOLATION",
88 "policy_reason": &message,
89 })
90 .to_string();
91
92 JobExecutionOutcome::FatalFailure {
93 message,
94 execution_id: None,
95 diagnostics: Some(diagnostics),
96 }
97 }
98
99}
100
101#[async_trait]
102impl JobHandler for PromptChatJobHandler {
103 fn job_type(&self) -> &'static str {
104 "workflow.stasis.prompt"
105 }
106
107 async fn execute(&self, job: &Job) -> Result<JobExecutionOutcome> {
108 let payload = match Self::parse_payload(&job.payload_ref) {
109 Ok(payload) => payload,
110 Err(message) => return Ok(Self::build_failure(message)),
111 };
112
113 let execution_context = RuntimeHandlerExecutionContext::new(
114 job,
115 payload.policy_profile.clone(),
116 payload.model_hint.clone(),
117 self.memory_reader.is_some(),
118 self.memory_writer.is_some(),
119 self.identity_memory_store.is_some(),
120 );
121
122 let mut memory_recall = None;
123 let mut memory_recall_error = None;
124 let mut input_memory_query_id = None;
125 let mut input_memory_query_fingerprint = None;
126
127 let memory_policy = payload.memory_policy.as_ref();
128 let (identity_summary, identity_error) = load_identity_context_summary(
129 self.identity_memory_store.as_ref(),
130 execution_context.correlation_id(),
131 execution_context.policy_profile(),
132 )
133 .await;
134
135 let mut effective_user_prompt =
136 prepend_identity_snapshot(&payload.user_prompt, identity_summary.as_deref());
137
138 if let Some(reader) = &self.memory_reader {
139 let recall_request = build_memory_recall_request(
140 execution_context.correlation_id(),
141 Some(&effective_user_prompt),
142 memory_policy,
143 );
144 input_memory_query_id = Some(memory_query_id(
145 execution_context.correlation_id(),
146 &recall_request,
147 ));
148 input_memory_query_fingerprint = Some(memory_query_fingerprint(&recall_request));
149
150 match reader.recall(&recall_request).await {
151 Ok(response) => {
152 effective_user_prompt = prepend_memory_recall_context(&effective_user_prompt, &response);
153 memory_recall = Some(response);
154 }
155 Err(err) => memory_recall_error = Some(err.to_string()),
156 }
157 }
158
159 let context = execution_context.prompt_context_clone();
160
161 let user_prompt = effective_user_prompt;
162 let mut request = PromptExecutionRequest::from_user_prompt(user_prompt.clone())
163 .with_context(context.clone());
164 if let Some(system_prompt) = payload.system_prompt {
165 request = request.with_system_prompt(system_prompt);
166 }
167
168 let response = match self.pipeline.execute(request).await {
169 Ok(response) => response,
170 Err(err) => {
171 let error = err.to_string();
172 return Ok(JobExecutionOutcome::FatalFailure {
173 message: error.clone(),
174 execution_id: None,
175 diagnostics: Some(
176 json!({
177 "provider": "stasis-pipeline",
178 "status": "failure",
179 "error": error,
180 "policy_profile": execution_context.policy_profile_clone(),
181 "model_hint": execution_context.model_hint_clone(),
182 "memory_recall": build_runtime_failure_memory_recall_section(
183 execution_context.memory_reader_enabled(),
184 memory_recall_error,
185 ),
186 "identity_context": build_runtime_failure_identity_context_section(
187 execution_context.identity_enabled(),
188 identity_summary,
189 identity_error,
190 ),
191 })
192 .to_string(),
193 ),
194 });
195 }
196 };
197
198 let mut memory_store = None;
199 let mut memory_store_error = None;
200 if should_store(memory_policy)
201 && let Some(writer) = &self.memory_writer
202 {
203 let store_request = MemoryStoreRequest {
204 session_id: execution_context.correlation_id().to_string(),
205 raw_node: render_prompt_response_sttp_node(
206 execution_context.correlation_id(),
207 &user_prompt,
208 &response.text,
209 SttpPromptNodeFormat::UntaggedNoSchema,
210 ),
211 };
212
213 match writer.store_context(&store_request).await {
214 Ok(stored) => memory_store = Some(stored),
215 Err(err) => memory_store_error = Some(err.to_string()),
216 }
217 }
218
219 let sttp_output_node_id =
220 resolve_sttp_output_node_id(memory_store.as_ref(), format!("sttp:prompt:{}", job.id));
221 let memory_scope_hash = memory_scope_hash(execution_context.correlation_id(), memory_policy);
222 let input_memory_query_id_for_top_level = input_memory_query_id.clone();
223 let input_memory_query_fingerprint_for_top_level =
224 input_memory_query_fingerprint.clone();
225 let diagnostics_bundle = build_runtime_memory_diagnostics_bundle(
226 RuntimeMemoryRecallDiagnosticsInput {
227 attempted: execution_context.memory_reader_enabled(),
228 response: memory_recall,
229 query_id: input_memory_query_id,
230 query_fingerprint: input_memory_query_fingerprint,
231 error: memory_recall_error,
232 },
233 RuntimeMemoryStoreDiagnosticsInput {
234 attempted: execution_context.memory_writer_enabled(),
235 response: memory_store,
236 error: memory_store_error,
237 },
238 RuntimeIdentityDiagnosticsInput {
239 attempted: execution_context.identity_enabled(),
240 summary: identity_summary,
241 error: identity_error,
242 },
243 );
244
245 let diagnostics = json!({
246 "provider": "stasis-pipeline",
247 "status": "success",
248 "trace_id": context.trace_id,
249 "correlation_id": context.correlation_id,
250 "policy_profile": response.metadata.policy_profile,
251 "model_hint": response.metadata.model_hint,
252 "output_text": response.text,
253 "output_preview": response.text.chars().take(160).collect::<String>(),
254 "memory_retrieved_count": diagnostics_bundle.retrieved_count,
255 "memory_retrieval_path": diagnostics_bundle.retrieval_path,
256 "memory_fallback_triggered": diagnostics_bundle.fallback_triggered,
257 "memory_fallback_reason": diagnostics_bundle.fallback_reason,
258 "memory_scope_hash": memory_scope_hash,
259 "memory_store_valid": diagnostics_bundle.store_valid,
260 "memory_store_node_id": diagnostics_bundle.store_node_id,
261 "input_memory_query_id": input_memory_query_id_for_top_level,
262 "input_memory_query_fingerprint": input_memory_query_fingerprint_for_top_level,
263 "output_memory_node_id": diagnostics_bundle.store_node_id,
264 "memory_recall": diagnostics_bundle.memory_recall,
265 "memory_store": diagnostics_bundle.memory_store,
266 "identity_context": diagnostics_bundle.identity_context,
267 })
268 .to_string();
269
270 Ok(JobExecutionOutcome::Success {
271 sttp_output_node_id,
272 execution_id: None,
273 diagnostics: Some(diagnostics),
274 })
275 }
276}