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