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