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