1use super::tenant::TenantContext;
18use super::trace::TraceContext;
19use crate::kernel::{CancellationPolicy, ExecutionId, ParentLink, ParentType, SpawnMode, StepId};
20use serde::{Deserialize, Serialize};
21use std::collections::HashMap;
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct SessionContext {
26 pub session_id: String,
28 pub created_at: chrono::DateTime<chrono::Utc>,
30 pub last_active_at: Option<chrono::DateTime<chrono::Utc>>,
32 pub metadata: HashMap<String, serde_json::Value>,
34}
35
36impl SessionContext {
37 pub fn new(session_id: impl Into<String>) -> Self {
39 Self {
40 session_id: session_id.into(),
41 created_at: chrono::Utc::now(),
42 last_active_at: None,
43 metadata: HashMap::new(),
44 }
45 }
46
47 pub fn touch(&mut self) {
49 self.last_active_at = Some(chrono::Utc::now());
50 }
51}
52
53#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct RuntimeContext {
74 pub execution_id: ExecutionId,
77 pub step_id: Option<StepId>,
79
80 pub parent: ParentLink,
83
84 pub tenant: TenantContext,
87
88 pub trace: TraceContext,
91
92 pub session: Option<SessionContext>,
95
96 pub created_at: chrono::DateTime<chrono::Utc>,
99
100 pub spawn_mode: Option<SpawnMode>,
104
105 pub cancellation_policy: CancellationPolicy,
107
108 pub parent_execution_id: Option<ExecutionId>,
110
111 pub metadata: HashMap<String, serde_json::Value>,
114}
115
116impl RuntimeContext {
117 pub fn new(execution_id: ExecutionId, parent: ParentLink, tenant: TenantContext) -> Self {
122 Self {
123 execution_id,
124 step_id: None,
125 parent,
126 tenant,
127 trace: TraceContext::new(),
128 session: None,
129 created_at: chrono::Utc::now(),
130 spawn_mode: None,
131 cancellation_policy: CancellationPolicy::default(),
132 parent_execution_id: None,
133 metadata: HashMap::new(),
134 }
135 }
136
137 pub fn from_user_message(
139 execution_id: ExecutionId,
140 message_id: impl Into<String>,
141 tenant: TenantContext,
142 ) -> Self {
143 Self::new(
144 execution_id,
145 ParentLink::from_user_message(message_id),
146 tenant,
147 )
148 }
149
150 pub fn with_step(mut self, step_id: StepId) -> Self {
154 self.step_id = Some(step_id);
155 self
156 }
157
158 pub fn with_trace(mut self, trace: TraceContext) -> Self {
160 self.trace = trace;
161 self
162 }
163
164 pub fn with_session(mut self, session: SessionContext) -> Self {
166 self.session = Some(session);
167 self
168 }
169
170 pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
172 self.metadata.insert(key.into(), value);
173 self
174 }
175
176 pub fn with_spawn_mode(mut self, spawn_mode: SpawnMode) -> Self {
178 self.spawn_mode = Some(spawn_mode);
179 self
180 }
181
182 pub fn with_cancellation_policy(mut self, policy: CancellationPolicy) -> Self {
184 self.cancellation_policy = policy;
185 self
186 }
187
188 pub fn child_context(&self, child_execution_id: ExecutionId, parent_step_id: &StepId) -> Self {
201 self.child_context_with_spawn_mode(
202 child_execution_id,
203 parent_step_id,
204 SpawnMode::child(false, false),
205 )
206 }
207
208 pub fn child_context_with_spawn_mode(
212 &self,
213 child_execution_id: ExecutionId,
214 parent_step_id: &StepId,
215 spawn_mode: SpawnMode,
216 ) -> Self {
217 Self {
218 execution_id: child_execution_id,
219 step_id: None,
220 parent: ParentLink::from_step(parent_step_id),
221 tenant: self.tenant.child_context(None),
222 trace: self.trace.child_span(),
223 session: self.session.clone(),
224 created_at: chrono::Utc::now(),
225 spawn_mode: Some(spawn_mode),
226 cancellation_policy: CancellationPolicy::default(),
227 parent_execution_id: Some(self.execution_id.clone()),
228 metadata: HashMap::new(), }
230 }
231
232 pub fn enter_step(&self, step_id: StepId) -> Self {
234 let mut ctx = self.clone();
235 ctx.step_id = Some(step_id);
236 ctx.trace = ctx.trace.child_span();
237 ctx
238 }
239
240 pub fn execution_id(&self) -> &ExecutionId {
244 &self.execution_id
245 }
246
247 pub fn step_id(&self) -> Option<&StepId> {
249 self.step_id.as_ref()
250 }
251
252 pub fn tenant(&self) -> &TenantContext {
254 &self.tenant
255 }
256
257 pub fn trace_id(&self) -> &str {
259 &self.trace.trace_id
260 }
261
262 pub fn span_id(&self) -> &str {
264 &self.trace.span_id
265 }
266
267 pub fn is_root(&self) -> bool {
269 !matches!(self.parent.parent_type, ParentType::StepExecution)
270 }
271}
272
273pub struct RuntimeContextBuilder {
275 execution_id: ExecutionId,
276 parent: ParentLink,
277 tenant: TenantContext,
278 trace: TraceContext,
279 session: Option<SessionContext>,
280 spawn_mode: Option<SpawnMode>,
281 cancellation_policy: CancellationPolicy,
282 parent_execution_id: Option<ExecutionId>,
283 metadata: HashMap<String, serde_json::Value>,
284}
285
286impl RuntimeContextBuilder {
287 pub fn new(execution_id: ExecutionId, parent: ParentLink, tenant: TenantContext) -> Self {
291 Self {
292 execution_id,
293 parent,
294 tenant,
295 trace: TraceContext::new(),
296 session: None,
297 spawn_mode: None,
298 cancellation_policy: CancellationPolicy::default(),
299 parent_execution_id: None,
300 metadata: HashMap::new(),
301 }
302 }
303
304 pub fn trace(mut self, trace: TraceContext) -> Self {
306 self.trace = trace;
307 self
308 }
309
310 pub fn session(mut self, session: SessionContext) -> Self {
312 self.session = Some(session);
313 self
314 }
315
316 pub fn metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
318 self.metadata.insert(key.into(), value);
319 self
320 }
321
322 pub fn spawn_mode(mut self, spawn_mode: SpawnMode) -> Self {
324 self.spawn_mode = Some(spawn_mode);
325 self
326 }
327
328 pub fn cancellation_policy(mut self, policy: CancellationPolicy) -> Self {
330 self.cancellation_policy = policy;
331 self
332 }
333
334 pub fn parent_execution_id(mut self, parent_exec_id: ExecutionId) -> Self {
336 self.parent_execution_id = Some(parent_exec_id);
337 self
338 }
339
340 pub fn build(self) -> RuntimeContext {
342 RuntimeContext {
343 execution_id: self.execution_id,
344 step_id: None,
345 parent: self.parent,
346 tenant: self.tenant,
347 trace: self.trace,
348 session: self.session,
349 created_at: chrono::Utc::now(),
350 spawn_mode: self.spawn_mode,
351 cancellation_policy: self.cancellation_policy,
352 parent_execution_id: self.parent_execution_id,
353 metadata: self.metadata,
354 }
355 }
356}
357
358#[cfg(test)]
359mod tests {
360 use super::*;
361 use crate::kernel::{TenantId, UserId};
362
363 #[test]
368 fn test_session_context_new() {
369 let session = SessionContext::new("sess_123");
370 assert_eq!(session.session_id, "sess_123");
371 assert!(session.last_active_at.is_none());
372 assert!(session.metadata.is_empty());
373 }
374
375 #[test]
376 fn test_session_context_new_owned_string() {
377 let session = SessionContext::new(String::from("sess_owned"));
378 assert_eq!(session.session_id, "sess_owned");
379 }
380
381 #[test]
382 fn test_session_context_touch() {
383 let mut session = SessionContext::new("sess_touch");
384 assert!(session.last_active_at.is_none());
385
386 session.touch();
387 assert!(session.last_active_at.is_some());
388
389 let first_touch = session.last_active_at.unwrap();
390 std::thread::sleep(std::time::Duration::from_millis(10));
391 session.touch();
392
393 assert!(session.last_active_at.unwrap() >= first_touch);
394 }
395
396 #[test]
397 fn test_session_context_metadata() {
398 let mut session = SessionContext::new("sess_meta");
399 session
400 .metadata
401 .insert("key".to_string(), serde_json::json!("value"));
402
403 assert_eq!(
404 session.metadata.get("key"),
405 Some(&serde_json::json!("value"))
406 );
407 }
408
409 #[test]
410 fn test_session_context_serde() {
411 let session = SessionContext::new("sess_serde");
412 let json = serde_json::to_string(&session).unwrap();
413 let parsed: SessionContext = serde_json::from_str(&json).unwrap();
414 assert_eq!(session.session_id, parsed.session_id);
415 }
416
417 fn create_test_tenant() -> TenantContext {
422 TenantContext::new(TenantId::from_string("tenant_test"))
423 }
424
425 #[test]
426 fn test_runtime_context_new() {
427 let exec_id = ExecutionId::from_string("exec_test");
428 let parent = ParentLink::from_user_message("msg_123");
429 let tenant = create_test_tenant();
430
431 let ctx = RuntimeContext::new(exec_id.clone(), parent, tenant);
432
433 assert_eq!(ctx.execution_id.as_str(), "exec_test");
434 assert!(ctx.step_id.is_none());
435 assert!(ctx.session.is_none());
436 assert!(ctx.metadata.is_empty());
437 }
438
439 #[test]
440 fn test_runtime_context_from_user_message() {
441 let exec_id = ExecutionId::from_string("exec_msg");
442 let tenant = create_test_tenant();
443
444 let ctx = RuntimeContext::from_user_message(exec_id, "msg_456", tenant);
445
446 assert_eq!(ctx.parent.parent_type, ParentType::UserMessage);
447 assert_eq!(ctx.parent.parent_id, "msg_456");
448 }
449
450 #[test]
451 fn test_runtime_context_with_step() {
452 let exec_id = ExecutionId::from_string("exec_step");
453 let parent = ParentLink::system();
454 let tenant = create_test_tenant();
455 let step_id = StepId::from_string("step_ctx");
456
457 let ctx = RuntimeContext::new(exec_id, parent, tenant).with_step(step_id.clone());
458
459 assert!(ctx.step_id.is_some());
460 assert_eq!(ctx.step_id.unwrap().as_str(), "step_ctx");
461 }
462
463 #[test]
464 fn test_runtime_context_with_trace() {
465 let exec_id = ExecutionId::from_string("exec_trace");
466 let parent = ParentLink::system();
467 let tenant = create_test_tenant();
468 let trace = TraceContext::from_traceparent(
469 "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01",
470 )
471 .unwrap();
472
473 let ctx = RuntimeContext::new(exec_id, parent, tenant).with_trace(trace);
474
475 assert_eq!(ctx.trace_id(), "0123456789abcdef0123456789abcdef");
476 }
477
478 #[test]
479 fn test_runtime_context_with_session() {
480 let exec_id = ExecutionId::from_string("exec_sess");
481 let parent = ParentLink::system();
482 let tenant = create_test_tenant();
483 let session = SessionContext::new("sess_runtime");
484
485 let ctx = RuntimeContext::new(exec_id, parent, tenant).with_session(session);
486
487 assert!(ctx.session.is_some());
488 assert_eq!(ctx.session.unwrap().session_id, "sess_runtime");
489 }
490
491 #[test]
492 fn test_runtime_context_with_metadata() {
493 let exec_id = ExecutionId::from_string("exec_meta");
494 let parent = ParentLink::system();
495 let tenant = create_test_tenant();
496
497 let ctx = RuntimeContext::new(exec_id, parent, tenant)
498 .with_metadata("key1", serde_json::json!("value1"))
499 .with_metadata("key2", serde_json::json!(42));
500
501 assert_eq!(ctx.metadata.len(), 2);
502 assert_eq!(ctx.metadata.get("key1"), Some(&serde_json::json!("value1")));
503 assert_eq!(ctx.metadata.get("key2"), Some(&serde_json::json!(42)));
504 }
505
506 #[test]
507 fn test_runtime_context_child_context() {
508 let exec_id = ExecutionId::from_string("exec_parent");
509 let parent = ParentLink::from_user_message("msg_parent");
510 let tenant = TenantContext::new(TenantId::from_string("tenant_parent"))
511 .with_user(UserId::from_string("user_parent"));
512
513 let parent_ctx = RuntimeContext::new(exec_id, parent, tenant)
514 .with_session(SessionContext::new("sess_inherit"))
515 .with_metadata("parent_key", serde_json::json!("should_not_inherit"));
516
517 let child_exec_id = ExecutionId::from_string("exec_child");
518 let parent_step_id = StepId::from_string("step_that_spawned");
519 let child_ctx = parent_ctx.child_context(child_exec_id.clone(), &parent_step_id);
520
521 assert_eq!(child_ctx.execution_id.as_str(), "exec_child");
523
524 assert_eq!(child_ctx.parent.parent_type, ParentType::StepExecution);
526 assert_eq!(child_ctx.parent.parent_id, "step_that_spawned");
527
528 assert_eq!(child_ctx.tenant.tenant_id().as_str(), "tenant_parent");
530
531 assert!(child_ctx.session.is_some());
533 assert_eq!(child_ctx.session.unwrap().session_id, "sess_inherit");
534
535 assert_eq!(child_ctx.trace.trace_id, parent_ctx.trace.trace_id);
537 assert_ne!(child_ctx.trace.span_id, parent_ctx.trace.span_id);
538
539 assert!(child_ctx.metadata.is_empty());
541 }
542
543 #[test]
544 fn test_runtime_context_enter_step() {
545 let exec_id = ExecutionId::from_string("exec_enter");
546 let parent = ParentLink::system();
547 let tenant = create_test_tenant();
548 let original_ctx = RuntimeContext::new(exec_id, parent, tenant);
549
550 let step_id = StepId::from_string("step_enter");
551 let step_ctx = original_ctx.enter_step(step_id.clone());
552
553 assert!(step_ctx.step_id.is_some());
555 assert_eq!(step_ctx.step_id.unwrap().as_str(), "step_enter");
556
557 assert_eq!(step_ctx.trace.trace_id, original_ctx.trace.trace_id);
559 assert_ne!(step_ctx.trace.span_id, original_ctx.trace.span_id);
560
561 assert!(original_ctx.step_id.is_none());
563 }
564
565 #[test]
566 fn test_runtime_context_accessors() {
567 let exec_id = ExecutionId::from_string("exec_access");
568 let step_id = StepId::from_string("step_access");
569 let parent = ParentLink::system();
570 let tenant = TenantContext::new(TenantId::from_string("tenant_access"))
571 .with_user(UserId::from_string("user_access"));
572
573 let ctx = RuntimeContext::new(exec_id, parent, tenant).with_step(step_id);
574
575 assert_eq!(ctx.execution_id().as_str(), "exec_access");
576 assert_eq!(ctx.step_id().unwrap().as_str(), "step_access");
577 assert_eq!(ctx.tenant().tenant_id().as_str(), "tenant_access");
578 assert!(!ctx.trace_id().is_empty());
579 assert!(!ctx.span_id().is_empty());
580 }
581
582 #[test]
583 fn test_runtime_context_is_root_user_message() {
584 let exec_id = ExecutionId::from_string("exec_root");
585 let parent = ParentLink::from_user_message("msg_root");
586 let tenant = create_test_tenant();
587
588 let ctx = RuntimeContext::new(exec_id, parent, tenant);
589 assert!(ctx.is_root());
590 }
591
592 #[test]
593 fn test_runtime_context_is_root_system() {
594 let exec_id = ExecutionId::from_string("exec_root");
595 let parent = ParentLink::system();
596 let tenant = create_test_tenant();
597
598 let ctx = RuntimeContext::new(exec_id, parent, tenant);
599 assert!(ctx.is_root());
600 }
601
602 #[test]
603 fn test_runtime_context_is_not_root_step_execution() {
604 let exec_id = ExecutionId::from_string("exec_child");
605 let parent_step = StepId::from_string("step_parent");
606 let parent = ParentLink::from_step(&parent_step);
607 let tenant = create_test_tenant();
608
609 let ctx = RuntimeContext::new(exec_id, parent, tenant);
610 assert!(!ctx.is_root());
611 }
612
613 #[test]
614 fn test_runtime_context_serde() {
615 let exec_id = ExecutionId::from_string("exec_serde");
616 let parent = ParentLink::system();
617 let tenant = create_test_tenant();
618
619 let ctx = RuntimeContext::new(exec_id, parent, tenant);
620 let json = serde_json::to_string(&ctx).unwrap();
621 let parsed: RuntimeContext = serde_json::from_str(&json).unwrap();
622
623 assert_eq!(ctx.execution_id.as_str(), parsed.execution_id.as_str());
624 }
625
626 #[test]
631 fn test_builder_new() {
632 let exec_id = ExecutionId::from_string("exec_builder");
633 let parent = ParentLink::system();
634 let tenant = create_test_tenant();
635
636 let builder = RuntimeContextBuilder::new(exec_id.clone(), parent, tenant);
637 let ctx = builder.build();
638
639 assert_eq!(ctx.execution_id.as_str(), "exec_builder");
640 }
641
642 #[test]
643 fn test_builder_with_trace() {
644 let exec_id = ExecutionId::from_string("exec_builder_trace");
645 let parent = ParentLink::system();
646 let tenant = create_test_tenant();
647 let trace = TraceContext::from_traceparent(
648 "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1-bbbbbbbbbbbbbb11-01",
649 )
650 .unwrap();
651
652 let ctx = RuntimeContextBuilder::new(exec_id, parent, tenant)
653 .trace(trace)
654 .build();
655
656 assert_eq!(ctx.trace.trace_id, "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1");
657 }
658
659 #[test]
660 fn test_builder_with_session() {
661 let exec_id = ExecutionId::from_string("exec_builder_sess");
662 let parent = ParentLink::system();
663 let tenant = create_test_tenant();
664 let session = SessionContext::new("sess_builder");
665
666 let ctx = RuntimeContextBuilder::new(exec_id, parent, tenant)
667 .session(session)
668 .build();
669
670 assert!(ctx.session.is_some());
671 assert_eq!(ctx.session.unwrap().session_id, "sess_builder");
672 }
673
674 #[test]
675 fn test_builder_with_metadata() {
676 let exec_id = ExecutionId::from_string("exec_builder_meta");
677 let parent = ParentLink::system();
678 let tenant = create_test_tenant();
679
680 let ctx = RuntimeContextBuilder::new(exec_id, parent, tenant)
681 .metadata("build_key", serde_json::json!("build_value"))
682 .build();
683
684 assert_eq!(
685 ctx.metadata.get("build_key"),
686 Some(&serde_json::json!("build_value"))
687 );
688 }
689
690 #[test]
691 fn test_builder_full_chain() {
692 let exec_id = ExecutionId::from_string("exec_full");
693 let parent = ParentLink::from_user_message("msg_full");
694 let tenant = TenantContext::new(TenantId::from_string("tenant_full"))
695 .with_user(UserId::from_string("user_full"));
696 let session = SessionContext::new("sess_full");
697 let trace = TraceContext::new();
698
699 let ctx = RuntimeContextBuilder::new(exec_id, parent, tenant)
700 .trace(trace.clone())
701 .session(session)
702 .metadata("key1", serde_json::json!(1))
703 .metadata("key2", serde_json::json!(2))
704 .build();
705
706 assert_eq!(ctx.execution_id.as_str(), "exec_full");
707 assert_eq!(ctx.parent.parent_id, "msg_full");
708 assert_eq!(ctx.tenant.tenant_id().as_str(), "tenant_full");
709 assert!(ctx.session.is_some());
710 assert_eq!(ctx.metadata.len(), 2);
711 }
712
713 #[test]
718 fn test_nested_execution_hierarchy() {
719 let root_exec_id = ExecutionId::from_string("exec_root");
721 let tenant = TenantContext::new(TenantId::from_string("tenant_hier"))
722 .with_user(UserId::from_string("user_hier"));
723 let root_ctx = RuntimeContext::from_user_message(root_exec_id, "msg_root", tenant);
724
725 assert!(root_ctx.is_root());
726
727 let step1_id = StepId::from_string("step_1");
729 let child1_ctx = root_ctx.child_context(ExecutionId::from_string("exec_child1"), &step1_id);
730
731 assert!(!child1_ctx.is_root());
732 assert_eq!(child1_ctx.trace.trace_id, root_ctx.trace.trace_id);
733
734 let step2_id = StepId::from_string("step_2");
736 let child2_ctx =
737 child1_ctx.child_context(ExecutionId::from_string("exec_child2"), &step2_id);
738
739 assert!(!child2_ctx.is_root());
740 assert_eq!(child2_ctx.trace.trace_id, root_ctx.trace.trace_id);
742 assert_ne!(child2_ctx.trace.span_id, child1_ctx.trace.span_id);
744 assert_ne!(child1_ctx.trace.span_id, root_ctx.trace.span_id);
745 }
746
747 #[test]
748 fn test_step_execution_creates_correct_spans() {
749 let exec_id = ExecutionId::from_string("exec_spans");
750 let tenant = create_test_tenant();
751 let root_ctx = RuntimeContext::new(exec_id, ParentLink::system(), tenant);
752
753 let step1 = StepId::from_string("step_a");
754 let step2 = StepId::from_string("step_b");
755
756 let ctx_step1 = root_ctx.enter_step(step1.clone());
757 let ctx_step2 = root_ctx.enter_step(step2.clone());
758
759 assert_eq!(ctx_step1.trace.trace_id, ctx_step2.trace.trace_id);
761 assert_ne!(ctx_step1.trace.span_id, ctx_step2.trace.span_id);
763 assert_ne!(
765 ctx_step1.step_id.unwrap().as_str(),
766 ctx_step2.step_id.unwrap().as_str()
767 );
768 }
769}