1use std::sync::Arc;
2
3use async_trait::async_trait;
4use chrono::Utc;
5use serde_json::{Value, json};
6
7use crate::application::orchestration::runtime_job_payloads::{
8 AgentToolCallMode, ConcurrentBranchExecutionMode, ConcurrentBranchJobPayload,
9 ConcurrentPatternJobPayload, MemoryPolicyPayload,
10};
11use crate::application::orchestration::concurrent_pattern_pipeline::{
12 ConcurrentPatternBranch, ConcurrentPatternExecutionRequest, ConcurrentPatternPipeline,
13};
14use crate::application::orchestration::prompt_pipeline::PromptExecutionPipeline;
15use crate::application::orchestration::tool_loop_pipeline::ToolCallMode;
16use crate::application::orchestration::tool_registry::ToolRegistry;
17use crate::application::runtime::chat_options_resolver::validate_reasoning_effort;
18use crate::application::runtime::in_memory_runtime::{JobExecutionOutcome, JobHandler};
19use crate::domain::errors::Result;
20use crate::domain::runtime::job::Job;
21use crate::domain::runtime::thread::{NewThread, NewThreadEvent, ThreadMergeMetadata};
22use crate::ports::outbound::ai_chat_client::AiChatClient;
23use crate::ports::outbound::memory::identity_memory_store::IdentityMemoryStore;
24use crate::ports::outbound::memory::memory_context_reader::MemoryContextReader;
25use crate::ports::outbound::memory::memory_context_writer::MemoryContextWriter;
26use crate::ports::outbound::runtime::thread_store::ThreadStore;
27
28pub struct ConcurrentPatternJobHandler {
29 pipeline: ConcurrentPatternPipeline,
30 thread_store: Option<Arc<dyn ThreadStore>>,
31}
32
33impl ConcurrentPatternJobHandler {
34 pub fn new(chat_client: Arc<dyn AiChatClient>, tool_registry: Arc<dyn ToolRegistry>) -> Self {
35 Self::new_with_thread_store_and_memory(
36 chat_client,
37 tool_registry,
38 None,
39 None,
40 None,
41 None,
42 )
43 }
44
45 pub fn new_with_thread_store(
46 chat_client: Arc<dyn AiChatClient>,
47 tool_registry: Arc<dyn ToolRegistry>,
48 thread_store: Option<Arc<dyn ThreadStore>>,
49 ) -> Self {
50 Self::new_with_thread_store_and_memory(
51 chat_client,
52 tool_registry,
53 thread_store,
54 None,
55 None,
56 None,
57 )
58 }
59
60 pub fn new_with_thread_store_and_memory(
61 chat_client: Arc<dyn AiChatClient>,
62 tool_registry: Arc<dyn ToolRegistry>,
63 thread_store: Option<Arc<dyn ThreadStore>>,
64 memory_reader: Option<Arc<dyn MemoryContextReader>>,
65 memory_writer: Option<Arc<dyn MemoryContextWriter>>,
66 identity_memory_store: Option<Arc<dyn IdentityMemoryStore>>,
67 ) -> Self {
68 let prompt_pipeline = PromptExecutionPipeline::new(chat_client);
69 Self {
70 pipeline: ConcurrentPatternPipeline::new_with_tool_loop(
71 prompt_pipeline,
72 tool_registry,
73 memory_reader,
74 memory_writer,
75 identity_memory_store,
76 ),
77 thread_store,
78 }
79 }
80
81 async fn ensure_thread(
82 &self,
83 thread_id: &str,
84 parent_thread_id: Option<String>,
85 branch_label: Option<String>,
86 now: chrono::DateTime<Utc>,
87 ) {
88 let Some(store) = &self.thread_store else {
89 return;
90 };
91
92 let exists = store.get_thread(thread_id).await.ok().flatten().is_some();
93 if exists {
94 return;
95 }
96
97 let _ = store
98 .create_thread(NewThread {
99 thread_id: thread_id.to_string(),
100 parent_thread_id,
101 branch_label,
102 created_at: now,
103 })
104 .await;
105 }
106
107 async fn append_thread_event(
108 &self,
109 event_id: String,
110 thread_id: &str,
111 event_kind: &str,
112 payload_ref: String,
113 occurred_at: chrono::DateTime<Utc>,
114 ) {
115 let Some(store) = &self.thread_store else {
116 return;
117 };
118
119 let _ = store
120 .append_event(NewThreadEvent {
121 event_id,
122 thread_id: thread_id.to_string(),
123 event_kind: event_kind.to_string(),
124 payload_ref,
125 occurred_at,
126 })
127 .await;
128 }
129
130 fn parse_payload(raw: &str) -> std::result::Result<ConcurrentPatternJobPayload, String> {
131 let payload: ConcurrentPatternJobPayload = serde_json::from_str(raw).map_err(|err| {
132 format!("policy violation: invalid concurrent-pattern payload json: {err}")
133 })?;
134
135 if payload.initial_user_prompt.trim().is_empty() {
136 return Err(
137 "policy violation: concurrent-pattern payload.initial_user_prompt must be non-empty"
138 .to_string(),
139 );
140 }
141 if payload.branches.is_empty() {
142 return Err(
143 "policy violation: concurrent-pattern payload.branches must include at least one branch"
144 .to_string(),
145 );
146 }
147
148 for branch in &payload.branches {
149 Self::validate_branch(branch)?;
150 }
151
152 validate_reasoning_effort(payload.reasoning_effort.as_deref())
153 .map_err(|err| format!("policy violation: {err}"))?;
154
155 Ok(payload)
156 }
157
158 fn validate_branch(branch: &ConcurrentBranchJobPayload) -> std::result::Result<(), String> {
159 if branch.branch_id.trim().is_empty() {
160 return Err(
161 "policy violation: concurrent-pattern payload.branches[].branch_id must be non-empty"
162 .to_string(),
163 );
164 }
165 if branch.user_prompt_template.trim().is_empty() {
166 return Err(
167 "policy violation: concurrent-pattern payload.branches[].user_prompt_template must be non-empty"
168 .to_string(),
169 );
170 }
171 if branch.execution_mode == ConcurrentBranchExecutionMode::ToolLoop {
172 let tool_name = branch.tool_name.as_deref().unwrap_or_default().trim();
173 if tool_name.is_empty() {
174 return Err(
175 "policy violation: concurrent-pattern payload.branches[].tool_name must be non-empty when execution_mode is tool_loop"
176 .to_string(),
177 );
178 }
179 }
180
181 validate_reasoning_effort(branch.reasoning_effort.as_deref())
182 .map_err(|err| format!("policy violation: {err}"))?;
183
184 Ok(())
185 }
186
187 fn resolve_tool_call_mode(
188 branch_mode: Option<AgentToolCallMode>,
189 default_mode: Option<AgentToolCallMode>,
190 ) -> ToolCallMode {
191 match branch_mode.or(default_mode) {
192 Some(AgentToolCallMode::Strict) => ToolCallMode::Strict,
193 _ => ToolCallMode::Auto,
194 }
195 }
196
197 fn resolve_memory_policy(
198 branch_policy: Option<MemoryPolicyPayload>,
199 default_policy: Option<MemoryPolicyPayload>,
200 ) -> Option<MemoryPolicyPayload> {
201 branch_policy.or(default_policy)
202 }
203
204 fn build_failure(message: String) -> JobExecutionOutcome {
205 let diagnostics = json!({
206 "provider": "stasis-orchestration-concurrent",
207 "status": "failure",
208 "pattern": "concurrent",
209 "guardrail_code": "POLICY_VIOLATION",
210 "policy_reason": &message,
211 })
212 .to_string();
213
214 JobExecutionOutcome::FatalFailure {
215 message,
216 execution_id: None,
217 diagnostics: Some(diagnostics),
218 }
219 }
220}
221
222#[async_trait]
223impl JobHandler for ConcurrentPatternJobHandler {
224 fn job_type(&self) -> &'static str {
225 "workflow.stasis.orchestration.concurrent"
226 }
227
228 async fn execute(&self, job: &Job) -> Result<JobExecutionOutcome> {
229 let payload = match Self::parse_payload(&job.payload_ref) {
230 Ok(payload) => payload,
231 Err(message) => return Ok(Self::build_failure(message)),
232 };
233
234 let ConcurrentPatternJobPayload {
235 thread_id,
236 initial_user_prompt,
237 policy_profile,
238 model_hint,
239 reasoning_effort,
240 merge_strategy,
241 tool_call_mode,
242 memory_policy,
243 branches,
244 } = payload;
245
246 let pattern_tool_call_mode = tool_call_mode;
247 let pattern_memory_policy = memory_policy;
248
249 let now = Utc::now();
250 let thread_id = thread_id.unwrap_or_else(|| job.correlation_id.clone());
251 self.ensure_thread(&thread_id, None, Some("concurrent".to_string()), now)
252 .await;
253 self.append_thread_event(
254 format!("{}:concurrent:start", job.id),
255 &thread_id,
256 "orchestration.concurrent.started",
257 initial_user_prompt.clone(),
258 now,
259 )
260 .await;
261
262 let request = ConcurrentPatternExecutionRequest {
263 initial_user_prompt,
264 trace_id: Some(job.trace_id.clone()),
265 correlation_id: Some(job.correlation_id.clone()),
266 policy_profile,
267 model_hint,
268 reasoning_effort,
269 default_memory_policy: pattern_memory_policy.clone(),
270 merge_strategy,
271 branches: branches
272 .into_iter()
273 .map(|branch| {
274 let ConcurrentBranchJobPayload {
275 branch_id,
276 user_prompt_template,
277 system_prompt,
278 policy_profile,
279 model_hint,
280 reasoning_effort,
281 execution_mode,
282 tool_name,
283 tool_input,
284 tool_call_mode,
285 memory_policy,
286 } = branch;
287
288 ConcurrentPatternBranch {
289 branch_id,
290 user_prompt_template,
291 system_prompt,
292 policy_profile,
293 model_hint,
294 reasoning_effort,
295 execution_mode,
296 tool_name,
297 tool_input,
298 tool_call_mode: Self::resolve_tool_call_mode(
299 tool_call_mode,
300 pattern_tool_call_mode.clone(),
301 ),
302 memory_policy: Self::resolve_memory_policy(
303 memory_policy,
304 pattern_memory_policy.clone(),
305 ),
306 }
307 })
308 .collect(),
309 };
310
311 let response = match self.pipeline.execute(request).await {
312 Ok(response) => response,
313 Err(err) => {
314 let error = err.to_string();
315 return Ok(JobExecutionOutcome::FatalFailure {
316 message: error.clone(),
317 execution_id: None,
318 diagnostics: Some(
319 json!({
320 "provider": "stasis-orchestration-concurrent",
321 "status": "failure",
322 "pattern": "concurrent",
323 "error": error,
324 })
325 .to_string(),
326 ),
327 });
328 }
329 };
330
331 let branch_ids: Vec<String> = response
332 .branches
333 .iter()
334 .map(|branch| branch.branch_id.clone())
335 .collect();
336 let tool_loop_branch_count = response
337 .branches
338 .iter()
339 .filter(|branch| branch.execution_mode == ConcurrentBranchExecutionMode::ToolLoop)
340 .count();
341 let prompt_branch_count = response.branches.len() - tool_loop_branch_count;
342 let branch_summaries: Vec<Value> = response
343 .branches
344 .iter()
345 .map(|branch| {
346 json!({
347 "branch_id": branch.branch_id,
348 "execution_mode": match branch.execution_mode {
349 ConcurrentBranchExecutionMode::Prompt => "prompt",
350 ConcurrentBranchExecutionMode::ToolLoop => "tool_loop",
351 },
352 "rounds_executed": branch.rounds_executed,
353 "tool_invocation_count": branch.tool_invocations.len(),
354 "branch_termination_reason": branch.branch_termination_reason,
355 "memory_retrieved_count": branch.memory_retrieved_count,
356 "memory_store_node_id": branch.memory_store_node_id,
357 "input_memory_query_id": branch.input_memory_query_id,
358 "input_memory_query_fingerprint": branch.input_memory_query_fingerprint,
359 "memory_recall_error": branch.memory_recall_error,
360 "memory_store_error": branch.memory_store_error,
361 })
362 })
363 .collect();
364
365 let mut branch_thread_ids = Vec::new();
366 for branch in &response.branches {
367 let branch_thread_id = format!("{}::branch::{}", thread_id, branch.branch_id);
368 let branch_now = Utc::now();
369 self.ensure_thread(
370 &branch_thread_id,
371 Some(thread_id.clone()),
372 Some(branch.branch_id.clone()),
373 branch_now,
374 )
375 .await;
376 self.append_thread_event(
377 format!("{}:concurrent:branch:{}", job.id, branch.branch_id),
378 &branch_thread_id,
379 "orchestration.concurrent.branch.completed",
380 branch.output_text.clone(),
381 branch_now,
382 )
383 .await;
384 branch_thread_ids.push(branch_thread_id);
385 }
386
387 let merge_metadata = ThreadMergeMetadata {
388 parent_thread_id: thread_id.clone(),
389 branch_thread_ids: branch_thread_ids.clone(),
390 merge_strategy: response.merge_strategy.clone(),
391 merged_at: Utc::now(),
392 };
393 let merge_payload_ref =
394 serde_json::to_string(&merge_metadata).unwrap_or_else(|_| response.final_text.clone());
395
396 self.append_thread_event(
397 format!("{}:concurrent:completed", job.id),
398 &thread_id,
399 "orchestration.concurrent.completed",
400 merge_payload_ref,
401 Utc::now(),
402 )
403 .await;
404
405 Ok(JobExecutionOutcome::Success {
406 sttp_output_node_id: format!("sttp:orchestration:concurrent:{}", job.id),
407 execution_id: None,
408 diagnostics: Some(
409 json!({
410 "provider": "stasis-orchestration-concurrent",
411 "status": "success",
412 "pattern": "concurrent",
413 "branches_executed": response.branches.len(),
414 "prompt_branch_count": prompt_branch_count,
415 "tool_loop_branch_count": tool_loop_branch_count,
416 "branch_summaries": branch_summaries,
417 "branch_ids": branch_ids,
418 "thread_id": thread_id,
419 "branch_thread_ids": branch_thread_ids,
420 "thread_merge": merge_metadata,
421 "merge_strategy": response.merge_strategy,
422 "final_text": response.final_text,
423 "termination_reason": response.termination_reason,
424 })
425 .to_string(),
426 ),
427 })
428 }
429}