1use std::{
2 collections::HashMap,
3 sync::{
4 atomic::{AtomicU64, Ordering},
5 Arc, Mutex,
6 },
7 time::Duration,
8};
9
10use base64::Engine as _;
11use serde_json::Value;
12use sha2::{Digest, Sha256};
13use uuid::Uuid;
14
15use crate::{
16 adk::runtime_client::RuntimeServiceClient,
17 error::{Result, SdkError},
18 pb::{ActivationKind, ActivationRecoveryPolicy, BeginActivationRequest, WorkerSuspension},
19 runtime_adapter::ActivationDecision,
20};
21
22use super::config::ContextConfig;
23use super::registry::{FunctionCall, InvocationContext};
24
25#[derive(Debug, Clone)]
26pub struct CoreContext {
27 state: Arc<ContextState>,
28}
29
30#[derive(Debug)]
31struct ContextState {
32 client: Option<Arc<RuntimeServiceClient>>,
33 config: ContextConfig,
34 sleep_ordinal: AtomicU64,
35 completed_steps: Mutex<serde_json::Map<String, Value>>,
36}
37
38impl ContextState {
39 fn function_registry(&self) -> Arc<super::registry::FunctionRegistry> {
40 Arc::clone(&self.config.function_registry)
41 }
42}
43
44impl CoreContext {
45 pub fn new(client: Option<Arc<RuntimeServiceClient>>, config: ContextConfig) -> Self {
46 let completed_steps = completed_steps(&config.metadata);
47 let state = ContextState {
48 client,
49 config,
50 sleep_ordinal: AtomicU64::new(0),
51 completed_steps: Mutex::new(completed_steps),
52 };
53 Self {
54 state: Arc::new(state),
55 }
56 }
57
58 pub fn with_runtime(client: Arc<RuntimeServiceClient>, config: ContextConfig) -> Self {
59 Self::new(Some(client), config)
60 }
61
62 pub fn config(&self) -> &ContextConfig {
63 &self.state.config
64 }
65
66 pub fn runtime_client(&self) -> Option<&Arc<RuntimeServiceClient>> {
67 self.state.client.as_ref()
68 }
69
70 pub fn functions(&self) -> FunctionNamespace {
71 FunctionNamespace::new(self.state.clone())
72 }
73
74 pub fn signals(&self) -> SignalNamespace {
75 SignalNamespace::new(self.state.clone())
76 }
77
78 pub fn timers(&self) -> TimerNamespace {
79 TimerNamespace::new(self.state.clone())
80 }
81
82 pub fn language_model(&self) -> LanguageModelNamespace {
83 LanguageModelNamespace::new(self.state.clone())
84 }
85}
86
87#[derive(Debug, Clone)]
88pub struct FunctionNamespace {
89 state: Arc<ContextState>,
90}
91
92impl FunctionNamespace {
93 fn new(state: Arc<ContextState>) -> Self {
94 Self { state }
95 }
96
97 pub async fn call(&self, request: FunctionCall) -> Result<FunctionHandle> {
98 let registry = self.state.function_registry();
99 let invocation_id = Uuid::new_v4().to_string();
100 let invocation_ctx = InvocationContext::from(&self.state.config);
101 let original_request = request.clone();
102
103 match registry.invoke(request, invocation_ctx).await {
104 Ok(output) => Ok(FunctionHandle::succeeded(
105 original_request,
106 invocation_id,
107 output,
108 )),
109 Err(err) => match &err {
110 SdkError::InvalidArgument { .. } => Err(err),
111 _ => Ok(FunctionHandle::failed(
112 original_request,
113 invocation_id,
114 err.to_string(),
115 )),
116 },
117 }
118 }
119}
120
121#[derive(Debug, Clone)]
122pub struct FunctionHandle {
123 result: Arc<FunctionResult>,
124}
125
126impl FunctionHandle {
127 pub async fn result(&self) -> FunctionResult {
128 (*self.result).clone()
129 }
130
131 pub fn status(&self) -> FunctionStatus {
132 self.result.status
133 }
134
135 fn succeeded(request: FunctionCall, invocation_id: String, output: Value) -> Self {
136 let result = FunctionResult {
137 request,
138 invocation_id,
139 status: FunctionStatus::Succeeded,
140 output: Some(output),
141 error: None,
142 };
143 Self {
144 result: Arc::new(result),
145 }
146 }
147
148 fn failed(request: FunctionCall, invocation_id: String, error: String) -> Self {
149 let result = FunctionResult {
150 request,
151 invocation_id,
152 status: FunctionStatus::Failed,
153 output: None,
154 error: Some(error),
155 };
156 Self {
157 result: Arc::new(result),
158 }
159 }
160}
161
162#[derive(Debug, Clone, Copy, PartialEq, Eq)]
163pub enum FunctionStatus {
164 Pending,
165 Succeeded,
166 Failed,
167}
168
169#[derive(Debug, Clone)]
170pub struct FunctionResult {
171 pub request: FunctionCall,
172 pub invocation_id: String,
173 pub status: FunctionStatus,
174 pub output: Option<Value>,
175 pub error: Option<String>,
176}
177
178#[derive(Debug, Clone)]
179pub struct SignalNamespace {
180 state: Arc<ContextState>,
181}
182
183impl SignalNamespace {
184 fn new(state: Arc<ContextState>) -> Self {
185 Self { state }
186 }
187
188 pub async fn wait(&self, name: &str) -> Result<Value> {
189 let _ = (&self.state, name);
190 Err(SdkError::Unavailable {
191 message: "Signal waiting not yet implemented".to_string(),
192 service: None,
193 })
194 }
195
196 pub async fn emit(&self, name: &str, payload: Value) -> Result<()> {
197 let _ = (&self.state, name, &payload);
198 Err(SdkError::Unavailable {
199 message: "Signal emission not yet implemented".to_string(),
200 service: None,
201 })
202 }
203}
204
205#[derive(Debug, Clone)]
206pub struct TimerNamespace {
207 state: Arc<ContextState>,
208}
209
210impl TimerNamespace {
211 fn new(state: Arc<ContextState>) -> Self {
212 Self { state }
213 }
214
215 pub async fn sleep(&self, duration: Duration) -> Result<()> {
216 if duration.is_zero() {
217 return Ok(());
218 }
219 if self
220 .state
221 .config
222 .metadata
223 .get("durable_suspension_v1")
224 .map(String::as_str)
225 != Some("true")
226 {
227 return Err(activation_error(
228 crate::error::ErrorCode::DurabilityUnavailable,
229 "durable_suspension_v1 was not negotiated for this context",
230 None,
231 None,
232 ));
233 }
234
235 let ordinal = self.state.sleep_ordinal.fetch_add(1, Ordering::Relaxed);
236 let timer_key = format!("sleep:sleep_{ordinal}");
237 if self
238 .state
239 .completed_steps
240 .lock()
241 .unwrap_or_else(|poisoned| poisoned.into_inner())
242 .contains_key(&timer_key)
243 {
244 return Ok(());
245 }
246
247 let delay_ms = duration_to_delay_ms(duration)?;
248 let plan = durable_timer_plan(&self.state.config, timer_key.clone(), delay_ms)?;
249 if self
250 .state
251 .config
252 .metadata
253 .get("timer_key")
254 .is_some_and(|resumed| resumed == &timer_key)
255 {
256 let resumed_activation = self
257 .state
258 .config
259 .metadata
260 .get("activation_id")
261 .map(String::as_str)
262 .unwrap_or_default();
263 if resumed_activation != plan.expected_activation_id {
264 return Err(activation_error(
265 crate::error::ErrorCode::NondeterministicReplay,
266 "timer resume authority does not match the deterministic sleep activation",
267 (!resumed_activation.is_empty()).then(|| resumed_activation.to_string()),
268 None,
269 ));
270 }
271 self.state
272 .completed_steps
273 .lock()
274 .unwrap_or_else(|poisoned| poisoned.into_inner())
275 .insert(timer_key, Value::Null);
276 return Ok(());
277 }
278
279 let client = self
280 .state
281 .config
282 .timer_activation_client
283 .as_ref()
284 .ok_or_else(|| {
285 activation_error(
286 crate::error::ErrorCode::DurabilityUnavailable,
287 "durable timer activation client is unavailable",
288 None,
289 None,
290 )
291 })?;
292 let decision = client.begin_timer(plan.request).await?;
293 let execution = match decision {
294 ActivationDecision::Execute(receipt) => receipt,
295 decision => return Err(timer_decision_error(decision)),
296 };
297 let completed_steps = self
298 .state
299 .completed_steps
300 .lock()
301 .unwrap_or_else(|poisoned| poisoned.into_inner())
302 .clone();
303 let continuation = timer_continuation(&self.state.config.metadata, completed_steps)?;
304 Err(SdkError::DurableSuspension {
305 suspension: Box::new(WorkerSuspension {
306 activation_id: execution.activation_id,
307 attempt: execution.attempt,
308 fence_token: execution.fence_token,
309 timer_key,
310 ready_at_ms: 0,
311 input_digest: plan.input_digest,
312 definition_digest: plan.definition_digest,
313 continuation,
314 delay_ms,
315 }),
316 })
317 }
318}
319
320const ACTIVATION_IDENTITY_DOMAIN: &[u8] = b"agnt5.activation.identity.v1\0";
321const ACTIVATION_DEFINITION_DOMAIN: &[u8] = b"agnt5.activation.definition.v1\0";
322const DURABLE_ACTIVATION_V1: &[u8] = b"durable_activation_v1";
323
324struct DurableTimerPlan {
325 request: BeginActivationRequest,
326 expected_activation_id: String,
327 input_digest: Vec<u8>,
328 definition_digest: Vec<u8>,
329}
330
331fn durable_timer_plan(
332 config: &ContextConfig,
333 timer_key: String,
334 delay_ms: i64,
335) -> Result<DurableTimerPlan> {
336 let metadata = &config.metadata;
337 let project_id = metadata
338 .get("project_id")
339 .filter(|value| !value.is_empty())
340 .cloned()
341 .unwrap_or_else(|| config.tenant_id.clone());
342 let worker_session_id = metadata
343 .get("worker_session_id")
344 .or_else(|| metadata.get("worker_id"))
345 .cloned()
346 .unwrap_or_default();
347 let run_authority = metadata
348 .get("run_authority")
349 .cloned()
350 .or_else(|| config.invocation_id.clone())
351 .unwrap_or_else(|| config.run_id.clone());
352 let lease_authority = metadata
353 .get("lease_authority")
354 .or_else(|| metadata.get("lease_id"))
355 .cloned()
356 .unwrap_or_default();
357 let component_name = metadata
358 .get("component_name")
359 .cloned()
360 .unwrap_or_else(|| config.step_id.clone());
361 let definition_version = metadata
362 .get("activation_definition_version")
363 .cloned()
364 .unwrap_or_default();
365 if project_id.is_empty()
366 || config.run_id.is_empty()
367 || worker_session_id.is_empty()
368 || run_authority.is_empty()
369 || lease_authority.is_empty()
370 || component_name.is_empty()
371 || definition_version.is_empty()
372 {
373 return Err(activation_error(
374 crate::error::ErrorCode::DurabilityUnavailable,
375 "durable timer requires project, run, worker-session, run, lease, and definition authority",
376 None,
377 None,
378 ));
379 }
380
381 let canonical_input = serde_json::to_vec(&serde_json::json!([
382 "object",
383 [
384 ["delay_ms", ["i64", delay_ms.to_string()]],
385 ["timer_key", ["string", timer_key.clone()]],
386 ]
387 ]))?;
388 let input_digest = Sha256::digest(canonical_input).to_vec();
389 let artifact = decode_sha256(
390 metadata
391 .get("activation_artifact_sha256")
392 .map(String::as_str)
393 .unwrap_or_default(),
394 )?;
395 let canonical_config = metadata
396 .get("activation_definition_config")
397 .map(String::as_bytes)
398 .unwrap_or(b"[\"object\",[]]");
399 let parsed_config: Value = serde_json::from_slice(canonical_config).map_err(|error| {
400 activation_error(
401 crate::error::ErrorCode::InvalidInput,
402 format!("activation definition config is not valid canonical JSON: {error}"),
403 None,
404 None,
405 )
406 })?;
407 if serde_json::to_vec(&parsed_config)? != canonical_config {
408 return Err(activation_error(
409 crate::error::ErrorCode::InvalidInput,
410 "activation definition config is not canonically encoded",
411 None,
412 None,
413 ));
414 }
415 let mut definition_bytes = Vec::from(ACTIVATION_DEFINITION_DOMAIN);
416 for part in [
417 artifact.as_slice(),
418 component_name.as_bytes(),
419 definition_version.as_bytes(),
420 DURABLE_ACTIVATION_V1,
421 canonical_config,
422 ] {
423 push_frame(&mut definition_bytes, part);
424 }
425 let definition_digest = Sha256::digest(definition_bytes).to_vec();
426 let parent_activation_id = metadata
427 .get("parent_activation_id")
428 .cloned()
429 .unwrap_or_default();
430 let expected_activation_id = activation_id(
431 &project_id,
432 &config.run_id,
433 &parent_activation_id,
434 ActivationKind::Timer,
435 &timer_key,
436 );
437 let request = BeginActivationRequest {
438 project_id,
439 run_id: config.run_id.clone(),
440 parent_activation_id,
441 kind: ActivationKind::Timer as i32,
442 stable_key: timer_key.clone(),
443 input_digest: input_digest.clone(),
444 definition_digest: definition_digest.clone(),
445 recovery_policy: ActivationRecoveryPolicy::DurableSteps as i32,
446 worker_session_id,
447 run_authority: run_authority.into_bytes(),
448 lease_authority: lease_authority.into_bytes(),
449 child: None,
450 display_name: timer_key.clone(),
451 display_parent_correlation_id: String::new(),
452 input_data: serde_json::to_vec(&serde_json::json!({
453 "delay_ms": delay_ms,
454 "timer_key": timer_key,
455 }))?,
456 };
457 Ok(DurableTimerPlan {
458 request,
459 expected_activation_id,
460 input_digest,
461 definition_digest,
462 })
463}
464
465fn duration_to_delay_ms(duration: Duration) -> Result<i64> {
466 let nanos = duration.as_nanos();
467 let millis = nanos.saturating_add(999_999) / 1_000_000;
468 i64::try_from(millis).map_err(|_| SdkError::InvalidArgument {
469 message: "timer duration exceeds maximum supported range".to_string(),
470 argument: Some("duration".to_string()),
471 })
472}
473
474fn activation_id(
475 project_id: &str,
476 run_id: &str,
477 parent_activation_id: &str,
478 kind: ActivationKind,
479 stable_key: &str,
480) -> String {
481 let mut encoded = Vec::from(ACTIVATION_IDENTITY_DOMAIN);
482 for part in [
483 project_id.as_bytes(),
484 run_id.as_bytes(),
485 parent_activation_id.as_bytes(),
486 ] {
487 push_frame(&mut encoded, part);
488 }
489 encoded.extend_from_slice(&(kind as u32).to_be_bytes());
490 push_frame(&mut encoded, stable_key.as_bytes());
491 format!(
492 "actv1_{}",
493 base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(Sha256::digest(encoded))
494 )
495}
496
497fn push_frame(target: &mut Vec<u8>, value: &[u8]) {
498 target.extend_from_slice(&(value.len() as u64).to_be_bytes());
499 target.extend_from_slice(value);
500}
501
502fn decode_sha256(value: &str) -> Result<Vec<u8>> {
503 if value.is_empty() {
504 return Err(activation_error(
505 crate::error::ErrorCode::DurabilityUnavailable,
506 "activation artifact SHA-256 is unavailable",
507 None,
508 None,
509 ));
510 }
511 let decoded = hex::decode(value)
512 .ok()
513 .filter(|bytes| bytes.len() == 32)
514 .or_else(|| {
515 [
516 base64::engine::general_purpose::STANDARD.decode(value),
517 base64::engine::general_purpose::STANDARD_NO_PAD.decode(value),
518 base64::engine::general_purpose::URL_SAFE.decode(value),
519 base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(value),
520 ]
521 .into_iter()
522 .find_map(|decoded| decoded.ok().filter(|bytes| bytes.len() == 32))
523 });
524 decoded.ok_or_else(|| {
525 activation_error(
526 crate::error::ErrorCode::InvalidInput,
527 "activation artifact SHA-256 must encode exactly 32 bytes",
528 None,
529 None,
530 )
531 })
532}
533
534fn completed_steps(metadata: &HashMap<String, String>) -> serde_json::Map<String, Value> {
535 metadata
536 .get("completed_steps")
537 .and_then(|value| serde_json::from_str::<serde_json::Map<String, Value>>(value).ok())
538 .unwrap_or_default()
539}
540
541fn timer_continuation(
542 metadata: &HashMap<String, String>,
543 completed_steps: serde_json::Map<String, Value>,
544) -> Result<Vec<u8>> {
545 let mut continuation = serde_json::Map::new();
546 if !completed_steps.is_empty() {
547 continuation.insert(
548 "completed_steps".to_string(),
549 Value::Object(completed_steps),
550 );
551 }
552 for key in ["step_events", "workflow_state"] {
553 if let Some(value) = metadata.get(key) {
554 if let Ok(parsed) = serde_json::from_str(value) {
555 continuation.insert(key.to_string(), parsed);
556 }
557 }
558 }
559 if let Some(value) = metadata.get("workflow_correlation_id") {
560 continuation.insert(
561 "workflow_correlation_id".to_string(),
562 Value::String(value.clone()),
563 );
564 }
565 serde_json::to_vec(&Value::Object(continuation)).map_err(Into::into)
566}
567
568fn timer_decision_error(decision: ActivationDecision) -> SdkError {
569 match decision {
570 ActivationDecision::Wait { activation_id, .. } => activation_error(
571 crate::error::ErrorCode::ActivationContended,
572 "timer activation is already executing",
573 Some(activation_id),
574 None,
575 ),
576 ActivationDecision::Conflict {
577 activation_id,
578 receipt,
579 } => activation_error(
580 crate::error::ErrorCode::NondeterministicReplay,
581 receipt.message,
582 Some(activation_id),
583 None,
584 ),
585 ActivationDecision::Cancelled {
586 activation_id,
587 attempt,
588 ..
589 } => activation_error(
590 crate::error::ErrorCode::ActivationCancelled,
591 "timer activation was cancelled",
592 Some(activation_id),
593 Some(attempt),
594 ),
595 ActivationDecision::UnknownOutcome { activation_id, .. } => activation_error(
596 crate::error::ErrorCode::UnknownOutcome,
597 "timer activation has an unknown outcome",
598 Some(activation_id),
599 None,
600 ),
601 ActivationDecision::Replay(receipt) => activation_error(
602 crate::error::ErrorCode::UnknownOutcome,
603 "timer activation unexpectedly returned a terminal replay",
604 Some(receipt.activation_id),
605 Some(receipt.attempt),
606 ),
607 ActivationDecision::Execute(receipt) => activation_error(
608 crate::error::ErrorCode::UnknownOutcome,
609 "timer activation execution receipt was not consumed",
610 Some(receipt.activation_id),
611 Some(receipt.attempt),
612 ),
613 }
614}
615
616fn activation_error(
617 code: crate::error::ErrorCode,
618 message: impl Into<String>,
619 activation_id: Option<String>,
620 attempt: Option<u32>,
621) -> SdkError {
622 SdkError::Activation {
623 message: message.into(),
624 code,
625 activation_id,
626 attempt,
627 }
628}
629
630#[derive(Debug, Clone)]
631pub struct LanguageModelNamespace {
632 state: Arc<ContextState>,
633}
634
635impl LanguageModelNamespace {
636 fn new(state: Arc<ContextState>) -> Self {
637 Self { state }
638 }
639
640 pub async fn generate(&self, _request: serde_json::Value) -> Result<Value> {
641 let _ = &self.state;
642 Err(SdkError::Unavailable {
643 message: "Language model generation not yet implemented".to_string(),
644 service: None,
645 })
646 }
647}
648
649#[cfg(test)]
650mod tests {
651 use std::sync::{Arc, Mutex};
652 use std::time::Duration;
653
654 use serde_json::json;
655
656 use crate::context::config::ContextConfig;
657 use crate::context::registry::{FunctionCall, FunctionRegistry};
658 use crate::context::TimerActivationClient;
659 use crate::error::{Result, SdkError};
660 use crate::pb::{ActivationKind, BeginActivationRequest};
661 use crate::runtime_adapter::{ActivationDecision, ActivationExecutionReceipt};
662
663 use super::{activation_id, CoreContext, FunctionStatus};
664
665 #[derive(Debug)]
666 struct TimerActivationStub {
667 requests: Mutex<Vec<BeginActivationRequest>>,
668 decision: ActivationDecision,
669 }
670
671 #[async_trait::async_trait]
672 impl TimerActivationClient for TimerActivationStub {
673 async fn begin_timer(&self, request: BeginActivationRequest) -> Result<ActivationDecision> {
674 self.requests
675 .lock()
676 .unwrap_or_else(|poisoned| poisoned.into_inner())
677 .push(request);
678 Ok(self.decision.clone())
679 }
680 }
681
682 fn durable_timer_config() -> ContextConfig {
683 ContextConfig::new("project-1", "session-1", "run-1", "workflow", 0)
684 .with_invocation_id("run-1")
685 .with_metadata("durable_suspension_v1", "true")
686 .with_metadata("project_id", "project-1")
687 .with_metadata("worker_session_id", "worker-1")
688 .with_metadata("lease_authority", "lease-1")
689 .with_metadata("component_name", "workflow")
690 .with_metadata("activation_definition_version", "v1")
691 .with_metadata("activation_definition_config", "[\"object\",[]]")
692 .with_metadata("activation_artifact_sha256", "00".repeat(32))
693 }
694
695 #[test]
696 fn context_stores_configuration() {
697 let cfg =
698 ContextConfig::new("tenant", "session", "run", "step", 0).with_invocation_id("invoke");
699 let ctx = CoreContext::new(None, cfg.clone());
700
701 assert_eq!(ctx.config().invocation_id, cfg.invocation_id);
702 assert_eq!(ctx.config().tenant_id, cfg.tenant_id);
703 }
704
705 #[tokio::test]
706 async fn timer_zero_duration_is_an_immediate_noop() {
707 let ctx = CoreContext::new(None, ContextConfig::default());
708 ctx.timers()
709 .sleep(Duration::ZERO)
710 .await
711 .expect("zero timer");
712 }
713
714 #[tokio::test]
715 async fn timer_sleep_returns_typed_durable_suspension() {
716 let timer_key = "sleep:sleep_0";
717 let expected_id = activation_id("project-1", "run-1", "", ActivationKind::Timer, timer_key);
718 let stub = Arc::new(TimerActivationStub {
719 requests: Mutex::new(Vec::new()),
720 decision: ActivationDecision::Execute(ActivationExecutionReceipt {
721 activation_id: expected_id.clone(),
722 attempt: 1,
723 fence_token: b"fence-1".to_vec(),
724 accepted_journal_offset: 7,
725 }),
726 });
727 let config = durable_timer_config().with_timer_activation_client(stub.clone());
728 let ctx = CoreContext::new(None, config);
729
730 let error = ctx
731 .timers()
732 .sleep(Duration::from_millis(2_500))
733 .await
734 .expect_err("timer must suspend");
735 let SdkError::DurableSuspension { suspension } = error else {
736 panic!("expected typed suspension");
737 };
738 assert_eq!(suspension.activation_id, expected_id);
739 assert_eq!(
740 suspension.activation_id,
741 "actv1_D1ifjG2fzE7kuiPRcKCoqPqIOKFwcesV0ttRbcK3Og0"
742 );
743 assert_eq!(suspension.timer_key, timer_key);
744 assert_eq!(suspension.delay_ms, 2_500);
745 assert_eq!(suspension.fence_token, b"fence-1");
746 assert_eq!(suspension.input_digest.len(), 32);
747 assert_eq!(suspension.definition_digest.len(), 32);
748 assert_eq!(
749 hex::encode(&suspension.input_digest),
750 "d1b146bf89e149e9df5ead1964ea34dc77c0b4b9e6b9d2dc8d7e2c58f0501ee7"
751 );
752 assert_eq!(
753 hex::encode(&suspension.definition_digest),
754 "18372aa11ce6f79e1e218982c1bc22976cfc28f868ac4a1c13811e4416e973a6"
755 );
756
757 let requests = stub
758 .requests
759 .lock()
760 .unwrap_or_else(|poisoned| poisoned.into_inner());
761 assert_eq!(requests.len(), 1);
762 assert_eq!(requests[0].kind, ActivationKind::Timer as i32);
763 assert_eq!(requests[0].stable_key, timer_key);
764 assert_eq!(requests[0].run_authority, b"run-1");
765 assert_eq!(requests[0].lease_authority, b"lease-1");
766 assert_eq!(requests[0].display_name, timer_key);
767 let input_data: serde_json::Value =
768 serde_json::from_slice(&requests[0].input_data).expect("input_data is JSON");
769 assert_eq!(
770 input_data,
771 json!({"delay_ms": 2_500, "timer_key": timer_key})
772 );
773 }
774
775 #[tokio::test]
776 async fn timer_resume_requires_the_matching_deterministic_activation() {
777 let timer_key = "sleep:sleep_0";
778 let expected_id = activation_id("project-1", "run-1", "", ActivationKind::Timer, timer_key);
779 let config = durable_timer_config()
780 .with_metadata("timer_key", timer_key)
781 .with_metadata("activation_id", expected_id);
782 let ctx = CoreContext::new(None, config);
783
784 ctx.timers()
785 .sleep(Duration::from_secs(1))
786 .await
787 .expect("matching resume");
788 }
789
790 #[tokio::test]
791 async fn timer_replay_skips_completed_sleep_before_current_resume() {
792 let second_key = "sleep:sleep_1";
793 let second_id = activation_id("project-1", "run-1", "", ActivationKind::Timer, second_key);
794 let config = durable_timer_config()
795 .with_metadata("completed_steps", r#"{"sleep:sleep_0":null}"#)
796 .with_metadata("timer_key", second_key)
797 .with_metadata("activation_id", second_id);
798 let ctx = CoreContext::new(None, config);
799
800 ctx.timers()
801 .sleep(Duration::from_secs(1))
802 .await
803 .expect("completed first sleep");
804 ctx.timers()
805 .sleep(Duration::from_secs(1))
806 .await
807 .expect("current second sleep");
808 }
809
810 #[tokio::test]
811 async fn timer_continuation_includes_sleep_completed_during_replay() {
812 let first_key = "sleep:sleep_0";
813 let first_id = activation_id("project-1", "run-1", "", ActivationKind::Timer, first_key);
814 let second_key = "sleep:sleep_1";
815 let second_id = activation_id("project-1", "run-1", "", ActivationKind::Timer, second_key);
816 let stub = Arc::new(TimerActivationStub {
817 requests: Mutex::new(Vec::new()),
818 decision: ActivationDecision::Execute(ActivationExecutionReceipt {
819 activation_id: second_id,
820 attempt: 1,
821 fence_token: b"fence-2".to_vec(),
822 accepted_journal_offset: 9,
823 }),
824 });
825 let config = durable_timer_config()
826 .with_metadata("timer_key", first_key)
827 .with_metadata("activation_id", first_id)
828 .with_timer_activation_client(stub);
829 let ctx = CoreContext::new(None, config);
830
831 ctx.timers()
832 .sleep(Duration::from_secs(1))
833 .await
834 .expect("resume first sleep");
835 let error = ctx
836 .timers()
837 .sleep(Duration::from_secs(1))
838 .await
839 .expect_err("second sleep must suspend");
840 let SdkError::DurableSuspension { suspension } = error else {
841 panic!("expected typed suspension");
842 };
843 let continuation: serde_json::Value =
844 serde_json::from_slice(&suspension.continuation).expect("continuation json");
845 assert_eq!(
846 continuation["completed_steps"][first_key],
847 serde_json::Value::Null
848 );
849 }
850
851 #[tokio::test]
852 async fn function_namespace_invokes_registered_function() {
853 let registry = Arc::new(FunctionRegistry::new());
854 registry.register("analytics", "process", |call, ctx| async move {
855 assert_eq!(ctx.run_id, "run");
856 assert!(call.metadata.get("corr_id").is_some());
857 Ok(json!({
858 "echo": call.payload,
859 }))
860 });
861
862 let cfg = ContextConfig::new("tenant", "session", "run", "step", 0)
863 .with_function_registry(Arc::clone(®istry));
864 let ctx = CoreContext::new(None, cfg);
865 let request = FunctionCall::new("analytics", "process", json!({"foo": "bar"}))
866 .with_metadata("corr_id", "123");
867
868 let handle = ctx.functions().call(request).await.expect("call succeeds");
869 assert_eq!(handle.status(), FunctionStatus::Succeeded);
870
871 let result = handle.result().await;
872 assert_eq!(result.status, FunctionStatus::Succeeded);
873 assert_eq!(result.output.unwrap()["echo"]["foo"], "bar");
874 assert!(!result.invocation_id.is_empty());
875 }
876
877 #[tokio::test]
878 async fn function_namespace_handles_handler_error() {
879 let registry = Arc::new(FunctionRegistry::new());
880 registry.register("svc", "fail", |_call, _ctx| async move {
881 Err(crate::error::SdkError::Invocation {
882 message: "boom".into(),
883 function_name: None,
884 })
885 });
886
887 let cfg = ContextConfig::new("tenant", "session", "run", "step", 0)
888 .with_function_registry(Arc::clone(®istry));
889 let ctx = CoreContext::new(None, cfg);
890 let request = FunctionCall::new("svc", "fail", json!({}));
891
892 let handle = ctx.functions().call(request).await.expect("call succeeds");
893 let result = handle.result().await;
894 assert_eq!(result.status, FunctionStatus::Failed);
895 assert!(result.error.unwrap().contains("boom"));
896 }
897
898 #[tokio::test]
899 async fn function_namespace_errors_when_unregistered() {
900 let cfg = ContextConfig::new("tenant", "session", "run", "step", 0);
901 let ctx = CoreContext::new(None, cfg);
902 let request = FunctionCall::new("missing", "handler", json!({}));
903
904 let err = ctx.functions().call(request).await;
905 assert!(matches!(err, Err(SdkError::InvalidArgument { .. })));
906 }
907}