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