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