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