Skip to main content

durable_workflow/
lib.rs

1//! Minimal Rust SDK for the Durable Workflow worker protocol.
2//!
3//! The crate covers the v1 Rust round-trip: start, signal, query, durably sleep,
4//! and start and await child workflows; register a Rust worker; poll workflow,
5//! activity, and read-only query tasks; reconstruct typed workflow-instance
6//! state through deterministic replay; heartbeat worker and activity liveness;
7//! and exchange JSON-native payloads through the same `avro` generic wrapper
8//! used by the existing first-party SDKs.
9
10use std::{
11    any::{Any, TypeId},
12    collections::{BTreeMap, HashMap},
13    future::Future,
14    pin::Pin,
15    sync::{
16        atomic::{AtomicBool, Ordering},
17        Arc, Mutex, OnceLock,
18    },
19    task::{Context as TaskContext, Poll},
20    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
21};
22
23use apache_avro::{from_avro_datum, from_value, to_avro_datum, to_value, Schema};
24use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
25use futures_util::{future::OptionFuture, task::noop_waker_ref};
26use serde::{de::DeserializeOwned, Deserialize, Serialize};
27pub use serde_json::{json, Value};
28use thiserror::Error;
29
30pub const WORKER_PROTOCOL_VERSION: &str = "1.2";
31pub const CONTROL_PLANE_VERSION: &str = "2";
32pub const DEFAULT_CODEC: &str = "avro";
33pub const JSON_CODEC: &str = "json";
34pub const SDK_VERSION: &str = concat!("durable-workflow-rust/", env!("CARGO_PKG_VERSION"));
35/// Worker-registration capability for server-routed read-only queries.
36pub const QUERY_TASKS_CAPABILITY: &str = "query_tasks";
37/// First additive worker protocol that defines query-task transport.
38pub const QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION: &str = "1.8";
39
40const MAX_LONG_POLL_TIMEOUT_SECONDS: u64 = 60;
41
42const QUERY_TASK_FINAL_REJECTION_REASONS: &[&str] = &[
43    "lease_expired",
44    "query_task_not_found",
45    "query_task_not_leased",
46    "query_task_timed_out",
47];
48
49const AVRO_PAYLOAD_SCHEMA_JSON: &str = r#"{"type":"record","name":"Payload","namespace":"durable_workflow","fields":[{"name":"json","type":"string"},{"name":"version","type":"int","default":1}]}"#;
50const AVRO_PAYLOAD_VERSION: i32 = 1;
51
52static AVRO_PAYLOAD_SCHEMA: OnceLock<std::result::Result<Schema, String>> = OnceLock::new();
53
54#[derive(Clone, Copy)]
55enum RequestProtocol {
56    ControlPlane,
57    Worker(&'static str),
58}
59
60impl RequestProtocol {
61    fn is_worker(self) -> bool {
62        matches!(self, Self::Worker(_))
63    }
64}
65
66pub type Result<T> = std::result::Result<T, Error>;
67
68#[derive(Debug, Error)]
69pub enum Error {
70    #[error("transport error: {0}")]
71    Transport(#[from] reqwest::Error),
72    #[error("json error: {0}")]
73    Json(#[from] serde_json::Error),
74    #[error("http {status}: {body}")]
75    Http {
76        status: reqwest::StatusCode,
77        body: String,
78    },
79    #[error("codec error: {0}")]
80    Codec(String),
81    #[error(transparent)]
82    QueryFailed(QueryFailure),
83    #[error(transparent)]
84    Protocol(ProtocolFailure),
85    #[error(transparent)]
86    NonDeterministicReplay(ReplayFailure),
87    #[error(transparent)]
88    ChildWorkflowFailed(ChildWorkflowFailure),
89    #[error("workflow handler {0:?} is not registered")]
90    WorkflowNotRegistered(String),
91    #[error("activity handler {0:?} is not registered")]
92    ActivityNotRegistered(String),
93    #[error("workflow future yielded without emitting a durable command")]
94    WorkflowYieldedWithoutCommand,
95    #[error("workflow state lock is poisoned")]
96    WorkflowStatePoisoned,
97    #[error("timer duration is too large for the worker protocol")]
98    TimerDurationOverflow,
99    #[error("operation timed out")]
100    Timeout,
101    #[error("worker loop error: {0}")]
102    WorkerLoop(String),
103    #[error("invalid child workflow options: {0}")]
104    InvalidChildWorkflowOptions(String),
105}
106
107/// Stable terminal categories returned when an awaited child does not succeed.
108#[derive(Clone, Copy, Debug, PartialEq, Eq)]
109pub enum ChildWorkflowFailureKind {
110    Failed,
111    Cancelled,
112    Terminated,
113}
114
115/// A stable, machine-readable child workflow failure delivered to its parent.
116///
117/// Match [`Error::ChildWorkflowFailed`] and inspect `reason` or `kind` instead
118/// of parsing the display message. Child and parent identifiers retain the
119/// relationship recorded in durable history across worker restarts.
120#[derive(Clone, Debug, Error)]
121#[error("child workflow failed ({reason}): {message}")]
122pub struct ChildWorkflowFailure {
123    pub kind: ChildWorkflowFailureKind,
124    pub reason: String,
125    pub message: String,
126    pub parent_workflow_id: Option<String>,
127    pub parent_workflow_run_id: Option<String>,
128    pub child_workflow_id: Option<String>,
129    pub child_workflow_run_id: Option<String>,
130    pub child_workflow_type: Option<String>,
131    pub failure_id: Option<String>,
132    pub failure_category: Option<String>,
133    pub exception_type: Option<String>,
134    pub exception_class: Option<String>,
135    pub non_retryable: bool,
136    pub code: Option<Value>,
137    pub exception: Option<Value>,
138}
139
140/// The identity of one durable workflow execution.
141#[derive(Clone, Debug, PartialEq, Eq)]
142pub struct WorkflowIdentity {
143    pub workflow_id: Option<String>,
144    pub run_id: Option<String>,
145}
146
147/// A successful child result together with its durable parent-child identity.
148#[derive(Clone, Debug, PartialEq)]
149pub struct ChildWorkflowResult {
150    pub parent: WorkflowIdentity,
151    pub child: WorkflowIdentity,
152    pub child_workflow_type: Option<String>,
153    pub result: Value,
154}
155
156/// Server behavior when a parent closes while its child is still open.
157#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
158pub enum ParentClosePolicy {
159    #[default]
160    Abandon,
161    RequestCancel,
162    Terminate,
163}
164
165impl ParentClosePolicy {
166    fn as_str(self) -> &'static str {
167        match self {
168            Self::Abandon => "abandon",
169            Self::RequestCancel => "request_cancel",
170            Self::Terminate => "terminate",
171        }
172    }
173}
174
175/// Durable retry policy for one child workflow invocation.
176#[derive(Clone, Debug, Default, PartialEq, Eq)]
177pub struct ChildWorkflowRetryPolicy {
178    pub max_attempts: Option<u32>,
179    pub backoff_seconds: Vec<u64>,
180    pub non_retryable_error_types: Vec<String>,
181}
182
183/// Options recorded with a child-workflow command.
184///
185/// The task queue is mandatory so routing is explicit and replay-stable.
186#[derive(Clone, Debug, PartialEq, Eq)]
187pub struct ChildWorkflowOptions {
188    pub task_queue: String,
189    pub parent_close_policy: ParentClosePolicy,
190    pub retry_policy: Option<ChildWorkflowRetryPolicy>,
191    pub execution_timeout_seconds: Option<u64>,
192    pub run_timeout_seconds: Option<u64>,
193}
194
195impl ChildWorkflowOptions {
196    pub fn new(task_queue: impl Into<String>) -> Self {
197        Self {
198            task_queue: task_queue.into(),
199            parent_close_policy: ParentClosePolicy::Abandon,
200            retry_policy: None,
201            execution_timeout_seconds: None,
202            run_timeout_seconds: None,
203        }
204    }
205
206    pub fn parent_close_policy(mut self, policy: ParentClosePolicy) -> Self {
207        self.parent_close_policy = policy;
208        self
209    }
210
211    pub fn retry_policy(mut self, policy: ChildWorkflowRetryPolicy) -> Self {
212        self.retry_policy = Some(policy);
213        self
214    }
215
216    pub fn execution_timeout_seconds(mut self, seconds: u64) -> Self {
217        self.execution_timeout_seconds = Some(seconds);
218        self
219    }
220
221    pub fn run_timeout_seconds(mut self, seconds: u64) -> Self {
222        self.run_timeout_seconds = Some(seconds);
223        self
224    }
225}
226
227/// A stable, machine-readable failure raised when workflow code no longer
228/// reconstructs the durable command stream recorded in history.
229#[derive(Clone, Debug, Error)]
230#[error("non-deterministic workflow replay ({reason}) at sequence {sequence:?}: {message}")]
231pub struct ReplayFailure {
232    pub reason: String,
233    pub sequence: Option<u64>,
234    pub expected: Option<String>,
235    pub actual: Option<String>,
236    pub message: String,
237}
238
239impl ReplayFailure {
240    fn new(
241        reason: impl Into<String>,
242        sequence: Option<u64>,
243        expected: Option<String>,
244        actual: Option<String>,
245        message: impl Into<String>,
246    ) -> Self {
247        Self {
248            reason: reason.into(),
249            sequence,
250            expected,
251            actual,
252            message: message.into(),
253        }
254    }
255}
256
257/// A stable, machine-readable workflow query or query-task settlement failure.
258#[derive(Clone, Debug, Error)]
259#[error("query failed ({reason}, HTTP {status}): {message}")]
260pub struct QueryFailure {
261    pub status: u16,
262    pub reason: String,
263    pub message: String,
264    pub body: Value,
265}
266
267/// A stable failure returned when a server rejects an SDK protocol version.
268#[derive(Clone, Debug, Error)]
269#[error("protocol rejected ({reason}, HTTP {status}): {message}")]
270pub struct ProtocolFailure {
271    pub status: u16,
272    pub reason: String,
273    pub message: String,
274    pub supported_version: Option<String>,
275    pub requested_version: Option<String>,
276    pub body: Value,
277}
278
279#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
280pub struct PayloadEnvelope {
281    pub codec: String,
282    pub blob: String,
283}
284
285impl PayloadEnvelope {
286    pub fn avro<T: Serialize>(value: &T) -> Result<Self> {
287        encode_payload(value, DEFAULT_CODEC)
288    }
289
290    pub fn json<T: Serialize>(value: &T) -> Result<Self> {
291        encode_payload(value, JSON_CODEC)
292    }
293}
294
295pub fn encode_payload<T: Serialize>(value: &T, codec: &str) -> Result<PayloadEnvelope> {
296    let value = serde_json::to_value(value)?;
297    let blob = encode_value_blob(&value, codec)?;
298
299    Ok(PayloadEnvelope {
300        codec: codec.to_string(),
301        blob,
302    })
303}
304
305pub fn decode_payload<T: DeserializeOwned>(envelope: &PayloadEnvelope) -> Result<T> {
306    let value = decode_blob(&envelope.blob, &envelope.codec)?;
307    Ok(serde_json::from_value(value)?)
308}
309
310fn encode_value_envelope(value: &Value, codec: &str) -> Result<Value> {
311    Ok(serde_json::to_value(encode_payload(value, codec)?)?)
312}
313
314fn encode_value_blob(value: &Value, codec: &str) -> Result<String> {
315    match codec {
316        JSON_CODEC => Ok(serde_json::to_string(value)?),
317        DEFAULT_CODEC => encode_avro_generic(value),
318        other => Err(Error::Codec(format!("unsupported payload codec {other:?}"))),
319    }
320}
321
322fn decode_wire_value(value: &Value, fallback_codec: &str) -> Result<Value> {
323    if value.is_null() {
324        return Ok(Value::Null);
325    }
326
327    if let Some(object) = value.as_object() {
328        if let (Some(codec), Some(blob)) = (
329            object.get("codec").and_then(Value::as_str),
330            object.get("blob").and_then(Value::as_str),
331        ) {
332            return decode_blob(blob, codec);
333        }
334    }
335
336    if let Some(blob) = value.as_str() {
337        return decode_blob(blob, fallback_codec);
338    }
339
340    Ok(value.clone())
341}
342
343fn decode_blob(blob: &str, codec: &str) -> Result<Value> {
344    match codec {
345        JSON_CODEC => Ok(serde_json::from_str(blob)?),
346        DEFAULT_CODEC => decode_avro_generic(blob),
347        other => Err(Error::Codec(format!("unsupported payload codec {other:?}"))),
348    }
349}
350
351fn encode_avro_generic(value: &Value) -> Result<String> {
352    let json = serde_json::to_string(value)?;
353    let datum = to_value(AvroPayload {
354        json,
355        version: AVRO_PAYLOAD_VERSION,
356    })
357    .map_err(|err| Error::Codec(format!("could not convert avro generic wrapper: {err}")))?;
358    let datum = to_avro_datum(avro_payload_schema()?, datum)
359        .map_err(|err| Error::Codec(format!("could not encode avro generic wrapper: {err}")))?;
360
361    let mut bytes = Vec::with_capacity(datum.len() + 1);
362    bytes.push(0x00);
363    bytes.extend_from_slice(&datum);
364    Ok(BASE64.encode(bytes))
365}
366
367fn decode_avro_generic(blob: &str) -> Result<Value> {
368    let bytes = BASE64
369        .decode(blob)
370        .map_err(|err| Error::Codec(format!("invalid avro base64 payload: {err}")))?;
371
372    if bytes.is_empty() {
373        return Err(Error::Codec("avro payload is empty".to_string()));
374    }
375
376    match bytes[0] {
377        0x00 => {}
378        0x01 => {
379            return Err(Error::Codec(
380                "typed avro payloads require a schema context; v1 supports the generic wrapper"
381                    .to_string(),
382            ));
383        }
384        other => {
385            return Err(Error::Codec(format!(
386                "unknown avro payload prefix 0x{other:02x}"
387            )));
388        }
389    }
390
391    let mut datum = &bytes[1..];
392    let datum = from_avro_datum(avro_payload_schema()?, &mut datum, None)
393        .map_err(|err| Error::Codec(format!("could not decode avro generic wrapper: {err}")))?;
394    let payload: AvroPayload = from_value(&datum)
395        .map_err(|err| Error::Codec(format!("invalid avro generic wrapper record: {err}")))?;
396
397    if payload.version != AVRO_PAYLOAD_VERSION {
398        return Err(Error::Codec(format!(
399            "unsupported avro generic wrapper version {}",
400            payload.version
401        )));
402    }
403
404    Ok(serde_json::from_str(&payload.json)?)
405}
406
407#[derive(Debug, Serialize, Deserialize)]
408struct AvroPayload {
409    json: String,
410    version: i32,
411}
412
413fn avro_payload_schema() -> Result<&'static Schema> {
414    match AVRO_PAYLOAD_SCHEMA.get_or_init(|| {
415        Schema::parse_str(AVRO_PAYLOAD_SCHEMA_JSON)
416            .map_err(|err| format!("could not parse avro payload schema: {err}"))
417    }) {
418        Ok(schema) => Ok(schema),
419        Err(message) => Err(Error::Codec(message.clone())),
420    }
421}
422
423#[derive(Clone, Debug)]
424pub struct Client {
425    http: reqwest::Client,
426    base_url: String,
427    token: Option<String>,
428    control_token: Option<String>,
429    worker_token: Option<String>,
430    namespace: String,
431}
432
433impl Client {
434    pub fn new(base_url: impl Into<String>) -> Result<Self> {
435        Self::builder(base_url).build()
436    }
437
438    pub fn builder(base_url: impl Into<String>) -> ClientBuilder {
439        ClientBuilder {
440            base_url: base_url.into(),
441            token: None,
442            control_token: None,
443            worker_token: None,
444            namespace: "default".to_string(),
445            timeout: Duration::from_secs(60),
446        }
447    }
448
449    pub async fn health(&self) -> Result<Value> {
450        self.request_json(
451            reqwest::Method::GET,
452            "/health",
453            RequestProtocol::ControlPlane,
454            Option::<&Value>::None,
455        )
456        .await
457    }
458
459    pub async fn cluster_info(&self) -> Result<Value> {
460        self.request_json(
461            reqwest::Method::GET,
462            "/cluster/info",
463            RequestProtocol::ControlPlane,
464            Option::<&Value>::None,
465        )
466        .await
467    }
468
469    pub async fn start_workflow<T: Serialize>(
470        &self,
471        workflow_type: &str,
472        task_queue: &str,
473        workflow_id: &str,
474        input: T,
475    ) -> Result<WorkflowHandle> {
476        let input = serde_json::to_value(input)?;
477        let input_envelope = encode_value_envelope(&normalize_arguments(input), DEFAULT_CODEC)?;
478        let body = json!({
479            "workflow_id": workflow_id,
480            "workflow_type": workflow_type,
481            "task_queue": task_queue,
482            "input": input_envelope,
483            "execution_timeout_seconds": 3600,
484            "run_timeout_seconds": 600
485        });
486
487        let data: Value = self
488            .request_json(
489                reqwest::Method::POST,
490                "/workflows",
491                RequestProtocol::ControlPlane,
492                Some(&body),
493            )
494            .await?;
495
496        Ok(WorkflowHandle {
497            client: self.clone(),
498            workflow_id: data
499                .get("workflow_id")
500                .and_then(Value::as_str)
501                .unwrap_or(workflow_id)
502                .to_string(),
503            run_id: data
504                .get("run_id")
505                .and_then(Value::as_str)
506                .map(str::to_string),
507            workflow_type: data
508                .get("workflow_type")
509                .and_then(Value::as_str)
510                .unwrap_or(workflow_type)
511                .to_string(),
512        })
513    }
514
515    pub async fn signal_workflow<T: Serialize>(
516        &self,
517        workflow_id: &str,
518        signal_name: &str,
519        input: T,
520    ) -> Result<Value> {
521        let input = serde_json::to_value(input)?;
522        let input_envelope = encode_value_envelope(&normalize_arguments(input), DEFAULT_CODEC)?;
523        let body = json!({
524            "input": input_envelope
525        });
526        let path = format!("/workflows/{workflow_id}/signal/{signal_name}");
527        self.request_json(
528            reqwest::Method::POST,
529            &path,
530            RequestProtocol::ControlPlane,
531            Some(&body),
532        )
533        .await
534    }
535
536    /// Execute a named, read-only query against a running or completed workflow.
537    ///
538    /// Arguments and results use the platform payload envelope. Server and
539    /// worker rejections are returned as [`Error::QueryFailed`] with a stable
540    /// reason, HTTP status, and original response body.
541    pub async fn query_workflow<T: Serialize>(
542        &self,
543        workflow_id: &str,
544        query_name: &str,
545        input: T,
546    ) -> Result<Value> {
547        let input = serde_json::to_value(input)?;
548        let input_envelope = encode_value_envelope(&normalize_arguments(input), DEFAULT_CODEC)?;
549        let body = json!({
550            "input": input_envelope
551        });
552        let path = format!("/workflows/{workflow_id}/query/{query_name}");
553        let response: Value = match self
554            .request_json(
555                reqwest::Method::POST,
556                &path,
557                RequestProtocol::ControlPlane,
558                Some(&body),
559            )
560            .await
561        {
562            Ok(response) => response,
563            Err(Error::Http { status, body }) => {
564                return Err(Error::QueryFailed(query_failure(status, body)));
565            }
566            Err(error) => return Err(error),
567        };
568
569        if let Some(envelope) = response
570            .get("result_envelope")
571            .filter(|envelope| !envelope.is_null())
572        {
573            return decode_wire_value(envelope, DEFAULT_CODEC);
574        }
575
576        Ok(response.get("result").cloned().unwrap_or(Value::Null))
577    }
578
579    pub async fn describe_workflow(&self, workflow_id: &str) -> Result<WorkflowDescription> {
580        let path = format!("/workflows/{workflow_id}");
581        let mut data: WorkflowDescription = self
582            .request_json(
583                reqwest::Method::GET,
584                &path,
585                RequestProtocol::ControlPlane,
586                Option::<&Value>::None,
587            )
588            .await?;
589        data.decode_payloads()?;
590        Ok(data)
591    }
592
593    pub async fn register_worker(
594        &self,
595        worker_id: &str,
596        task_queue: &str,
597        supported_workflow_types: Vec<String>,
598        supported_activity_types: Vec<String>,
599        max_concurrent_workflow_tasks: usize,
600        max_concurrent_activity_tasks: usize,
601    ) -> Result<RegisterWorkerResponse> {
602        self.register_worker_with_capabilities(
603            worker_id,
604            task_queue,
605            supported_workflow_types,
606            supported_activity_types,
607            max_concurrent_workflow_tasks,
608            max_concurrent_activity_tasks,
609            Vec::new(),
610        )
611        .await
612    }
613
614    /// Register a worker and explicitly advertise additive worker capabilities.
615    pub async fn register_worker_with_capabilities(
616        &self,
617        worker_id: &str,
618        task_queue: &str,
619        supported_workflow_types: Vec<String>,
620        supported_activity_types: Vec<String>,
621        max_concurrent_workflow_tasks: usize,
622        max_concurrent_activity_tasks: usize,
623        capabilities: Vec<String>,
624    ) -> Result<RegisterWorkerResponse> {
625        let body = json!({
626            "worker_id": worker_id,
627            "task_queue": task_queue,
628            "runtime": "rust",
629            "sdk_version": SDK_VERSION,
630            "supported_workflow_types": supported_workflow_types,
631            "supported_activity_types": supported_activity_types,
632            "capabilities": capabilities,
633            "max_concurrent_workflow_tasks": max_concurrent_workflow_tasks,
634            "max_concurrent_activity_tasks": max_concurrent_activity_tasks
635        });
636
637        self.request_json(
638            reqwest::Method::POST,
639            "/worker/register",
640            RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
641            Some(&body),
642        )
643        .await
644    }
645
646    /// Long-poll for an ephemeral, read-only workflow query task.
647    pub async fn poll_query_task(
648        &self,
649        worker_id: &str,
650        task_queue: &str,
651        timeout: Duration,
652    ) -> Result<Option<QueryTask>> {
653        let timeout_seconds = long_poll_timeout_seconds(timeout);
654        let body = json!({
655            "worker_id": worker_id,
656            "task_queue": task_queue,
657            "poll_request_id": unique_request_id("rust-query-poll"),
658            "timeout_seconds": timeout_seconds,
659        });
660        let data: PollQueryTaskResponse = self
661            .request_json_with_timeout(
662                reqwest::Method::POST,
663                "/worker/query-tasks/poll",
664                RequestProtocol::Worker(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION),
665                Some(&body),
666                timeout + Duration::from_secs(5),
667            )
668            .await?;
669        Ok(data.task)
670    }
671
672    /// Complete a query task without appending workflow history.
673    pub async fn complete_query_task(
674        &self,
675        query_task_id: &str,
676        lease_owner: &str,
677        query_task_attempt: u64,
678        result: Value,
679        codec: &str,
680    ) -> Result<Value> {
681        let result_envelope = encode_value_envelope(&result, codec)?;
682        self.complete_query_task_with_envelope(
683            query_task_id,
684            lease_owner,
685            query_task_attempt,
686            result,
687            result_envelope,
688        )
689        .await
690    }
691
692    async fn complete_query_task_with_envelope(
693        &self,
694        query_task_id: &str,
695        lease_owner: &str,
696        query_task_attempt: u64,
697        result: Value,
698        result_envelope: Value,
699    ) -> Result<Value> {
700        let body = json!({
701            "lease_owner": lease_owner,
702            "query_task_attempt": query_task_attempt,
703            "result": result,
704            "result_envelope": result_envelope,
705        });
706        let path = format!("/worker/query-tasks/{query_task_id}/complete");
707        let response = self
708            .request_json(
709                reqwest::Method::POST,
710                &path,
711                RequestProtocol::Worker(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION),
712                Some(&body),
713            )
714            .await;
715        query_task_response(response)
716    }
717
718    /// Report a stable machine-readable query-task failure.
719    pub async fn fail_query_task(
720        &self,
721        query_task_id: &str,
722        lease_owner: &str,
723        query_task_attempt: u64,
724        message: impl Into<String>,
725        reason: impl Into<String>,
726        failure_type: impl Into<String>,
727    ) -> Result<Value> {
728        let body = json!({
729            "lease_owner": lease_owner,
730            "query_task_attempt": query_task_attempt,
731            "failure": {
732                "message": message.into(),
733                "reason": reason.into(),
734                "type": failure_type.into(),
735            }
736        });
737        let path = format!("/worker/query-tasks/{query_task_id}/fail");
738        let response = self
739            .request_json(
740                reqwest::Method::POST,
741                &path,
742                RequestProtocol::Worker(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION),
743                Some(&body),
744            )
745            .await;
746        query_task_response(response)
747    }
748
749    pub async fn heartbeat_worker(
750        &self,
751        worker_id: &str,
752        workflow_available: usize,
753        activity_available: usize,
754    ) -> Result<Value> {
755        let body = json!({
756            "worker_id": worker_id,
757            "task_slots": {
758                "workflow_available": workflow_available,
759                "activity_available": activity_available
760            },
761            "process_metrics": {
762                "process_id": std::process::id(),
763                "process_uptime_seconds": 0
764            }
765        });
766
767        self.request_json(
768            reqwest::Method::POST,
769            "/worker/heartbeat",
770            RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
771            Some(&body),
772        )
773        .await
774    }
775
776    pub async fn poll_workflow_task(
777        &self,
778        worker_id: &str,
779        task_queue: &str,
780        timeout: Duration,
781    ) -> Result<Option<WorkflowTask>> {
782        Ok(self
783            .poll_workflow_task_response(worker_id, task_queue, timeout)
784            .await?
785            .task)
786    }
787
788    pub async fn poll_workflow_task_response(
789        &self,
790        worker_id: &str,
791        task_queue: &str,
792        timeout: Duration,
793    ) -> Result<PollWorkflowTaskResponse> {
794        let body = json!({
795            "worker_id": worker_id,
796            "task_queue": task_queue,
797            "timeout_seconds": long_poll_timeout_seconds(timeout),
798        });
799        let mut data: PollWorkflowTaskResponse = self
800            .request_json_with_timeout(
801                reqwest::Method::POST,
802                "/worker/workflow-tasks/poll",
803                RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
804                Some(&body),
805                timeout + Duration::from_secs(5),
806            )
807            .await?;
808
809        if let Some(task) = data.task.as_mut() {
810            self.fetch_remaining_workflow_history(worker_id, task)
811                .await?;
812        }
813
814        Ok(data)
815    }
816
817    async fn fetch_remaining_workflow_history(
818        &self,
819        worker_id: &str,
820        task: &mut WorkflowTask,
821    ) -> Result<()> {
822        let mut next_token = task.next_history_page_token.clone();
823
824        while let Some(token) = next_token.take().filter(|token| !token.is_empty()) {
825            let lease_owner = task
826                .lease_owner
827                .clone()
828                .unwrap_or_else(|| worker_id.to_string());
829            let page = self
830                .workflow_task_history_page(
831                    &task.task_id,
832                    &lease_owner,
833                    task.workflow_task_attempt,
834                    &token,
835                )
836                .await?;
837
838            task.append_history_page(page);
839
840            if task.next_history_page_token.as_deref() == Some(token.as_str()) {
841                return Err(Error::Codec(
842                    "workflow history pagination returned the same page token".to_string(),
843                ));
844            }
845
846            next_token = task.next_history_page_token.clone();
847        }
848
849        Ok(())
850    }
851
852    async fn workflow_task_history_page(
853        &self,
854        task_id: &str,
855        lease_owner: &str,
856        workflow_task_attempt: u64,
857        next_history_page_token: &str,
858    ) -> Result<WorkflowTaskHistoryPage> {
859        let body = json!({
860            "lease_owner": lease_owner,
861            "workflow_task_attempt": workflow_task_attempt,
862            "next_history_page_token": next_history_page_token
863        });
864        let path = format!("/worker/workflow-tasks/{task_id}/history");
865
866        self.request_json(
867            reqwest::Method::POST,
868            &path,
869            RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
870            Some(&body),
871        )
872        .await
873    }
874
875    pub async fn complete_workflow_task(
876        &self,
877        task_id: &str,
878        lease_owner: &str,
879        workflow_task_attempt: u64,
880        commands: Vec<Value>,
881    ) -> Result<Value> {
882        let body = json!({
883            "lease_owner": lease_owner,
884            "workflow_task_attempt": workflow_task_attempt,
885            "commands": commands
886        });
887        let path = format!("/worker/workflow-tasks/{task_id}/complete");
888        self.request_json(
889            reqwest::Method::POST,
890            &path,
891            RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
892            Some(&body),
893        )
894        .await
895    }
896
897    pub async fn fail_workflow_task(
898        &self,
899        task_id: &str,
900        lease_owner: &str,
901        workflow_task_attempt: u64,
902        message: impl Into<String>,
903    ) -> Result<Value> {
904        let body = json!({
905            "lease_owner": lease_owner,
906            "workflow_task_attempt": workflow_task_attempt,
907            "failure": {
908                "message": message.into(),
909                "type": "RustWorkflowTaskFailure"
910            }
911        });
912        let path = format!("/worker/workflow-tasks/{task_id}/fail");
913        self.request_json(
914            reqwest::Method::POST,
915            &path,
916            RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
917            Some(&body),
918        )
919        .await
920    }
921
922    pub async fn poll_activity_task(
923        &self,
924        worker_id: &str,
925        task_queue: &str,
926        timeout: Duration,
927    ) -> Result<Option<ActivityTask>> {
928        let body = json!({
929            "worker_id": worker_id,
930            "task_queue": task_queue,
931            "timeout_seconds": long_poll_timeout_seconds(timeout),
932        });
933        let data: PollActivityTaskResponse = self
934            .request_json_with_timeout(
935                reqwest::Method::POST,
936                "/worker/activity-tasks/poll",
937                RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
938                Some(&body),
939                timeout + Duration::from_secs(5),
940            )
941            .await?;
942        Ok(data.task)
943    }
944
945    pub async fn complete_activity_task(
946        &self,
947        task_id: &str,
948        activity_attempt_id: &str,
949        lease_owner: &str,
950        result: Value,
951        codec: &str,
952    ) -> Result<Value> {
953        let result = encode_value_envelope(&result, codec)?;
954        let body = json!({
955            "activity_attempt_id": activity_attempt_id,
956            "lease_owner": lease_owner,
957            "result": result
958        });
959        let path = format!("/worker/activity-tasks/{task_id}/complete");
960        self.request_json(
961            reqwest::Method::POST,
962            &path,
963            RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
964            Some(&body),
965        )
966        .await
967    }
968
969    pub async fn fail_activity_task(
970        &self,
971        task_id: &str,
972        activity_attempt_id: &str,
973        lease_owner: &str,
974        message: impl Into<String>,
975        non_retryable: bool,
976    ) -> Result<Value> {
977        let body = json!({
978            "activity_attempt_id": activity_attempt_id,
979            "lease_owner": lease_owner,
980            "failure": {
981                "message": message.into(),
982                "type": "RustActivityFailure",
983                "non_retryable": non_retryable
984            }
985        });
986        let path = format!("/worker/activity-tasks/{task_id}/fail");
987        self.request_json(
988            reqwest::Method::POST,
989            &path,
990            RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
991            Some(&body),
992        )
993        .await
994    }
995
996    pub async fn heartbeat_activity_task(
997        &self,
998        task_id: &str,
999        activity_attempt_id: &str,
1000        lease_owner: &str,
1001        details: Value,
1002    ) -> Result<ActivityHeartbeatResponse> {
1003        let body = json!({
1004            "activity_attempt_id": activity_attempt_id,
1005            "lease_owner": lease_owner,
1006            "details": details
1007        });
1008        let path = format!("/worker/activity-tasks/{task_id}/heartbeat");
1009        self.request_json(
1010            reqwest::Method::POST,
1011            &path,
1012            RequestProtocol::Worker(WORKER_PROTOCOL_VERSION),
1013            Some(&body),
1014        )
1015        .await
1016    }
1017
1018    async fn request_json<T: DeserializeOwned, B: Serialize + ?Sized>(
1019        &self,
1020        method: reqwest::Method,
1021        path: &str,
1022        protocol: RequestProtocol,
1023        body: Option<&B>,
1024    ) -> Result<T> {
1025        self.request_json_with_timeout(method, path, protocol, body, Duration::from_secs(60))
1026            .await
1027    }
1028
1029    async fn request_json_with_timeout<T: DeserializeOwned, B: Serialize + ?Sized>(
1030        &self,
1031        method: reqwest::Method,
1032        path: &str,
1033        protocol: RequestProtocol,
1034        body: Option<&B>,
1035        timeout: Duration,
1036    ) -> Result<T> {
1037        let mut request = self
1038            .http
1039            .request(method, format!("{}/api{}", self.base_url, path))
1040            .timeout(timeout)
1041            .header(reqwest::header::ACCEPT, "application/json")
1042            .header(reqwest::header::CONTENT_TYPE, "application/json")
1043            .header("X-Namespace", &self.namespace);
1044
1045        match protocol {
1046            RequestProtocol::Worker(version) => {
1047                request = request.header("X-Durable-Workflow-Protocol-Version", version);
1048            }
1049            RequestProtocol::ControlPlane => {
1050                request = request.header(
1051                    "X-Durable-Workflow-Control-Plane-Version",
1052                    CONTROL_PLANE_VERSION,
1053                );
1054            }
1055        }
1056
1057        if let Some(token) = self.auth_token(protocol.is_worker()) {
1058            request = request.bearer_auth(token);
1059        }
1060
1061        if let Some(body) = body {
1062            request = request.json(body);
1063        }
1064
1065        let response = request.send().await?;
1066        let status = response.status();
1067        let bytes = response.bytes().await?;
1068
1069        if !status.is_success() {
1070            let body = String::from_utf8_lossy(&bytes).to_string();
1071            if let Some(protocol) = protocol_failure(status, &body) {
1072                return Err(Error::Protocol(protocol));
1073            }
1074            return Err(Error::Http { status, body });
1075        }
1076
1077        if bytes.is_empty() {
1078            return Ok(serde_json::from_value(Value::Null)?);
1079        }
1080
1081        Ok(serde_json::from_slice(&bytes)?)
1082    }
1083
1084    fn auth_token(&self, worker: bool) -> Option<&str> {
1085        if worker {
1086            self.worker_token
1087                .as_deref()
1088                .or(self.token.as_deref())
1089                .or(self.control_token.as_deref())
1090        } else {
1091            self.control_token
1092                .as_deref()
1093                .or(self.token.as_deref())
1094                .or(self.worker_token.as_deref())
1095        }
1096    }
1097}
1098
1099fn query_failure(status: reqwest::StatusCode, raw_body: String) -> QueryFailure {
1100    let body = serde_json::from_str(&raw_body).unwrap_or_else(|_| json!({"message": raw_body}));
1101    let reason = body
1102        .get("reason")
1103        .and_then(Value::as_str)
1104        .unwrap_or("query_rejected")
1105        .to_string();
1106    let message = body
1107        .get("message")
1108        .or_else(|| body.get("error"))
1109        .and_then(Value::as_str)
1110        .unwrap_or("workflow query was rejected")
1111        .to_string();
1112
1113    QueryFailure {
1114        status: status.as_u16(),
1115        reason,
1116        message,
1117        body,
1118    }
1119}
1120
1121fn query_task_response(response: Result<Value>) -> Result<Value> {
1122    match response {
1123        Err(Error::Http { status, body }) => Err(Error::QueryFailed(query_failure(status, body))),
1124        response => response,
1125    }
1126}
1127
1128fn query_task_rejection_is_final(error: &Error) -> bool {
1129    matches!(
1130        error,
1131        Error::QueryFailed(failure)
1132            if QUERY_TASK_FINAL_REJECTION_REASONS.contains(&failure.reason.as_str())
1133    )
1134}
1135
1136fn protocol_failure(status: reqwest::StatusCode, raw_body: &str) -> Option<ProtocolFailure> {
1137    let body: Value = serde_json::from_str(raw_body).ok()?;
1138    let reason = body.get("reason")?.as_str()?;
1139    if !matches!(
1140        reason,
1141        "missing_protocol_version"
1142            | "unsupported_protocol_version"
1143            | "missing_control_plane_version"
1144            | "unsupported_control_plane_version"
1145    ) {
1146        return None;
1147    }
1148
1149    Some(ProtocolFailure {
1150        status: status.as_u16(),
1151        reason: reason.to_string(),
1152        message: body
1153            .get("message")
1154            .or_else(|| body.get("error"))
1155            .and_then(Value::as_str)
1156            .unwrap_or("protocol version rejected")
1157            .to_string(),
1158        supported_version: body
1159            .get("supported_version")
1160            .and_then(Value::as_str)
1161            .map(str::to_string),
1162        requested_version: body
1163            .get("requested_version")
1164            .and_then(Value::as_str)
1165            .map(str::to_string),
1166        body,
1167    })
1168}
1169
1170fn long_poll_timeout_seconds(timeout: Duration) -> u64 {
1171    timeout
1172        .as_secs()
1173        .saturating_add(u64::from(timeout.subsec_nanos() > 0))
1174        .min(MAX_LONG_POLL_TIMEOUT_SECONDS)
1175}
1176
1177fn worker_operation_is_retryable(error: &Error) -> bool {
1178    match error {
1179        Error::Transport(error) => {
1180            error.is_timeout() || error.is_connect() || error.is_request() || error.is_body()
1181        }
1182        Error::Http { status, .. } => {
1183            matches!(
1184                *status,
1185                reqwest::StatusCode::REQUEST_TIMEOUT | reqwest::StatusCode::TOO_MANY_REQUESTS
1186            ) || status.is_server_error()
1187        }
1188        _ => false,
1189    }
1190}
1191
1192fn worker_retry_delay(policy: WorkerRetryPolicy, retry: usize) -> Duration {
1193    let exponent = retry.saturating_sub(1).min(31) as u32;
1194    policy
1195        .initial_backoff
1196        .saturating_mul(1_u32 << exponent)
1197        .min(policy.max_backoff)
1198}
1199
1200#[derive(Debug)]
1201pub struct ClientBuilder {
1202    base_url: String,
1203    token: Option<String>,
1204    control_token: Option<String>,
1205    worker_token: Option<String>,
1206    namespace: String,
1207    timeout: Duration,
1208}
1209
1210impl ClientBuilder {
1211    pub fn token(mut self, token: Option<String>) -> Self {
1212        self.token = token;
1213        self
1214    }
1215
1216    pub fn control_token(mut self, token: Option<String>) -> Self {
1217        self.control_token = token;
1218        self
1219    }
1220
1221    pub fn worker_token(mut self, token: Option<String>) -> Self {
1222        self.worker_token = token;
1223        self
1224    }
1225
1226    pub fn namespace(mut self, namespace: impl Into<String>) -> Self {
1227        self.namespace = namespace.into();
1228        self
1229    }
1230
1231    pub fn timeout(mut self, timeout: Duration) -> Self {
1232        self.timeout = timeout;
1233        self
1234    }
1235
1236    pub fn build(self) -> Result<Client> {
1237        Ok(Client {
1238            http: reqwest::Client::builder().timeout(self.timeout).build()?,
1239            base_url: self.base_url.trim_end_matches('/').to_string(),
1240            token: self.token,
1241            control_token: self.control_token,
1242            worker_token: self.worker_token,
1243            namespace: self.namespace,
1244        })
1245    }
1246}
1247
1248#[derive(Clone, Debug)]
1249pub struct WorkflowHandle {
1250    client: Client,
1251    pub workflow_id: String,
1252    pub run_id: Option<String>,
1253    pub workflow_type: String,
1254}
1255
1256impl WorkflowHandle {
1257    pub async fn describe(&self) -> Result<WorkflowDescription> {
1258        self.client.describe_workflow(&self.workflow_id).await
1259    }
1260
1261    pub async fn signal<T: Serialize>(&self, signal_name: &str, input: T) -> Result<Value> {
1262        self.client
1263            .signal_workflow(&self.workflow_id, signal_name, input)
1264            .await
1265    }
1266
1267    /// Execute a named, read-only query against this workflow.
1268    pub async fn query<T: Serialize>(&self, query_name: &str, input: T) -> Result<Value> {
1269        self.client
1270            .query_workflow(&self.workflow_id, query_name, input)
1271            .await
1272    }
1273
1274    pub async fn result(&self, options: WorkflowResultOptions) -> Result<Value> {
1275        let started = Instant::now();
1276
1277        loop {
1278            let description = self.describe().await?;
1279            if description.is_completed() {
1280                return Ok(description.output.unwrap_or(Value::Null));
1281            }
1282
1283            if description.is_terminal() {
1284                return Err(Error::Codec(format!(
1285                    "workflow {} closed with status {:?}",
1286                    self.workflow_id, description.status
1287                )));
1288            }
1289
1290            if started.elapsed() >= options.timeout {
1291                return Err(Error::Timeout);
1292            }
1293
1294            tokio::time::sleep(options.poll_interval).await;
1295        }
1296    }
1297}
1298
1299#[derive(Clone, Copy, Debug)]
1300pub struct WorkflowResultOptions {
1301    pub poll_interval: Duration,
1302    pub timeout: Duration,
1303}
1304
1305impl Default for WorkflowResultOptions {
1306    fn default() -> Self {
1307        Self {
1308            poll_interval: Duration::from_millis(500),
1309            timeout: Duration::from_secs(30),
1310        }
1311    }
1312}
1313
1314#[derive(Clone, Debug, Deserialize)]
1315pub struct WorkflowDescription {
1316    pub workflow_id: Option<String>,
1317    pub run_id: Option<String>,
1318    pub workflow_type: Option<String>,
1319    pub status: Option<String>,
1320    #[serde(default)]
1321    pub output: Option<Value>,
1322    #[serde(default)]
1323    pub output_envelope: Option<Value>,
1324    #[serde(flatten)]
1325    pub raw: HashMap<String, Value>,
1326}
1327
1328impl WorkflowDescription {
1329    pub fn is_completed(&self) -> bool {
1330        matches!(self.status.as_deref(), Some("completed" | "Completed"))
1331    }
1332
1333    pub fn is_terminal(&self) -> bool {
1334        matches!(
1335            self.status.as_deref(),
1336            Some(
1337                "completed"
1338                    | "Completed"
1339                    | "failed"
1340                    | "Failed"
1341                    | "cancelled"
1342                    | "Cancelled"
1343                    | "terminated"
1344                    | "Terminated"
1345                    | "timed_out"
1346                    | "TimedOut",
1347            )
1348        )
1349    }
1350
1351    fn decode_payloads(&mut self) -> Result<()> {
1352        if let Some(envelope) = &self.output_envelope {
1353            self.output = Some(decode_wire_value(envelope, DEFAULT_CODEC)?);
1354        }
1355
1356        Ok(())
1357    }
1358}
1359
1360#[derive(Clone, Debug, Deserialize)]
1361pub struct RegisterWorkerResponse {
1362    pub worker_id: String,
1363    pub registered: bool,
1364    #[serde(default)]
1365    pub heartbeat_interval_seconds: Option<u64>,
1366    #[serde(default)]
1367    pub protocol_version: Option<String>,
1368    #[serde(default)]
1369    pub server_capabilities: Option<Value>,
1370}
1371
1372#[derive(Clone, Debug, Deserialize)]
1373pub struct PollWorkflowTaskResponse {
1374    #[serde(default)]
1375    pub task: Option<WorkflowTask>,
1376    #[serde(default)]
1377    pub poll_status: Option<String>,
1378    #[serde(default)]
1379    pub reason: Option<String>,
1380    #[serde(default)]
1381    pub protocol_version: Option<String>,
1382    #[serde(default)]
1383    pub server_capabilities: Option<Value>,
1384}
1385
1386#[derive(Clone, Debug, Deserialize)]
1387struct PollActivityTaskResponse {
1388    #[serde(default)]
1389    task: Option<ActivityTask>,
1390}
1391
1392#[derive(Clone, Debug, Deserialize)]
1393struct PollQueryTaskResponse {
1394    #[serde(default)]
1395    task: Option<QueryTask>,
1396}
1397
1398/// An ephemeral server-routed query task.
1399#[derive(Clone, Debug, Deserialize)]
1400pub struct QueryTask {
1401    pub query_task_id: String,
1402    #[serde(default = "default_workflow_task_attempt")]
1403    pub query_task_attempt: u64,
1404    #[serde(default)]
1405    pub lease_owner: Option<String>,
1406    #[serde(default)]
1407    pub workflow_id: Option<String>,
1408    #[serde(default)]
1409    pub run_id: Option<String>,
1410    pub workflow_type: String,
1411    pub query_name: String,
1412    #[serde(default = "default_payload_codec")]
1413    pub payload_codec: String,
1414    #[serde(default)]
1415    pub workflow_arguments: Option<Value>,
1416    #[serde(default)]
1417    pub query_arguments: Option<Value>,
1418    #[serde(default)]
1419    pub history_events: Vec<HistoryEvent>,
1420    #[serde(default)]
1421    pub history_export: Option<Value>,
1422    #[serde(default)]
1423    pub run_status: Option<String>,
1424}
1425
1426#[derive(Clone, Debug, Deserialize)]
1427pub struct WorkflowTask {
1428    pub task_id: String,
1429    #[serde(default)]
1430    pub workflow_id: Option<String>,
1431    #[serde(default)]
1432    pub run_id: Option<String>,
1433    pub workflow_type: String,
1434    #[serde(default = "default_payload_codec")]
1435    pub payload_codec: String,
1436    #[serde(default)]
1437    pub arguments: Option<Value>,
1438    #[serde(default)]
1439    pub history_events: Vec<HistoryEvent>,
1440    #[serde(default)]
1441    pub total_history_events: Option<u64>,
1442    #[serde(default)]
1443    pub next_history_page_token: Option<String>,
1444    #[serde(default = "default_workflow_task_attempt")]
1445    pub workflow_task_attempt: u64,
1446    #[serde(default)]
1447    pub workflow_signal_id: Option<String>,
1448    #[serde(default)]
1449    pub signal_name: Option<String>,
1450    #[serde(default)]
1451    pub signal_arguments: Option<Value>,
1452    #[serde(default)]
1453    pub lease_owner: Option<String>,
1454}
1455
1456impl WorkflowTask {
1457    fn append_history_page(&mut self, page: WorkflowTaskHistoryPage) {
1458        self.history_events.extend(page.history_events);
1459
1460        if page.total_history_events.is_some() {
1461            self.total_history_events = page.total_history_events;
1462        }
1463
1464        self.next_history_page_token = page
1465            .next_history_page_token
1466            .filter(|token| !token.is_empty());
1467    }
1468}
1469
1470#[derive(Clone, Debug, Deserialize)]
1471struct WorkflowTaskHistoryPage {
1472    #[serde(default)]
1473    history_events: Vec<HistoryEvent>,
1474    #[serde(default)]
1475    total_history_events: Option<u64>,
1476    #[serde(default)]
1477    next_history_page_token: Option<String>,
1478}
1479
1480#[derive(Clone, Debug, Deserialize)]
1481pub struct ActivityTask {
1482    pub task_id: String,
1483    #[serde(default)]
1484    pub activity_attempt_id: Option<String>,
1485    #[serde(default)]
1486    pub attempt_id: Option<String>,
1487    pub activity_type: String,
1488    #[serde(default = "default_payload_codec")]
1489    pub payload_codec: String,
1490    #[serde(default)]
1491    pub arguments: Option<Value>,
1492    #[serde(default = "default_attempt_number")]
1493    pub attempt_number: u64,
1494    #[serde(default)]
1495    pub lease_owner: Option<String>,
1496}
1497
1498#[derive(Clone, Debug, Deserialize)]
1499pub struct HistoryEvent {
1500    #[serde(alias = "type")]
1501    pub event_type: String,
1502    #[serde(default)]
1503    pub payload: Value,
1504    #[serde(flatten)]
1505    pub raw: HashMap<String, Value>,
1506}
1507
1508/// One decoded signal in the committed workflow-history snapshot.
1509#[derive(Clone, Debug, PartialEq)]
1510pub struct QuerySignal {
1511    pub id: Option<String>,
1512    pub name: String,
1513    pub arguments: Vec<Value>,
1514    pub workflow_sequence: Option<u64>,
1515}
1516
1517/// Immutable state supplied to a registered query handler.
1518///
1519/// This context intentionally exposes no activity, signal-wait, or command
1520/// APIs. Query handlers inspect committed history and return a value; query
1521/// completion does not append an event or advance deterministic execution.
1522#[derive(Clone, Debug)]
1523pub struct QueryContext {
1524    pub workflow_id: Option<String>,
1525    pub run_id: Option<String>,
1526    pub workflow_type: String,
1527    pub run_status: Option<String>,
1528    workflow_input: Value,
1529    history_events: Arc<Vec<HistoryEvent>>,
1530    signal_events: Arc<Vec<QuerySignal>>,
1531}
1532
1533impl QueryContext {
1534    /// The normalized argument list used to start the workflow.
1535    pub fn workflow_input(&self) -> &Value {
1536        &self.workflow_input
1537    }
1538
1539    /// The immutable committed history used for this query snapshot.
1540    pub fn history_events(&self) -> &[HistoryEvent] {
1541        self.history_events.as_slice()
1542    }
1543
1544    /// All decoded signals in committed workflow order.
1545    pub fn signal_events(&self) -> &[QuerySignal] {
1546        self.signal_events.as_slice()
1547    }
1548
1549    /// Decoded argument lists for each committed signal with `signal_name`.
1550    pub fn signals(&self, signal_name: &str) -> Vec<Vec<Value>> {
1551        self.signal_events
1552            .iter()
1553            .filter(|signal| signal.name == signal_name)
1554            .map(|signal| signal.arguments.clone())
1555            .collect()
1556    }
1557}
1558
1559#[derive(Clone, Debug, Deserialize)]
1560pub struct ActivityHeartbeatResponse {
1561    #[serde(default)]
1562    pub cancel_requested: bool,
1563    #[serde(default)]
1564    pub heartbeat_recorded: bool,
1565}
1566
1567fn default_payload_codec() -> String {
1568    DEFAULT_CODEC.to_string()
1569}
1570
1571fn default_workflow_task_attempt() -> u64 {
1572    1
1573}
1574
1575fn default_attempt_number() -> u64 {
1576    1
1577}
1578
1579type WorkflowFuture = Pin<Box<dyn Future<Output = Result<Value>> + Send + 'static>>;
1580type WorkflowHandler = Arc<dyn Fn(WorkflowContext, Value) -> WorkflowFuture + Send + Sync>;
1581type ErasedWorkflowState = Arc<dyn Any + Send + Sync>;
1582type WorkflowStateSnapshot = Arc<dyn Fn() -> Result<ErasedWorkflowState> + Send + Sync>;
1583type ReplayedWorkflowHandler =
1584    Arc<dyn Fn(WorkflowContext, Value) -> ReplayedWorkflowInvocation + Send + Sync>;
1585type ActivityFuture = Pin<Box<dyn Future<Output = Result<Value>> + Send + 'static>>;
1586type ActivityHandler = Arc<dyn Fn(ActivityContext, Value) -> ActivityFuture + Send + Sync>;
1587type QueryFuture = Pin<Box<dyn Future<Output = Result<Value>> + Send + 'static>>;
1588type QueryHandler = Arc<dyn Fn(QueryContext, Value) -> QueryFuture + Send + Sync>;
1589type ReplayedQueryHandler = Arc<
1590    dyn Fn(QueryContext, ErasedWorkflowState, Value) -> std::result::Result<QueryFuture, String>
1591        + Send
1592        + Sync,
1593>;
1594type WorkerHeartbeatObserver = Arc<dyn Fn(&WorkerHeartbeatObservation) + Send + Sync>;
1595
1596struct ReplayedWorkflowInvocation {
1597    future: WorkflowFuture,
1598    snapshot: WorkflowStateSnapshot,
1599}
1600
1601#[derive(Clone)]
1602struct RegisteredWorkflow {
1603    execute: WorkflowHandler,
1604    replay: Option<ReplayedWorkflowHandler>,
1605    state_type: Option<TypeId>,
1606}
1607
1608#[derive(Clone)]
1609enum RegisteredQuery {
1610    Snapshot(QueryHandler),
1611    Replayed {
1612        state_type: TypeId,
1613        handler: ReplayedQueryHandler,
1614    },
1615}
1616
1617#[derive(Clone, Debug)]
1618pub struct WorkerHeartbeatObservation {
1619    pub worker_id: String,
1620    pub task_queue: String,
1621    pub acknowledged_at_unix_millis: u64,
1622    pub acknowledgement: Value,
1623}
1624
1625/// Bounded retry policy for worker poll acquisition and worker heartbeats.
1626///
1627/// Expected empty long polls are normal successful responses. Transport
1628/// failures, HTTP 408/429 responses, and server errors are retried with capped
1629/// exponential backoff. Authentication, protocol, codec, and handler failures
1630/// are never retried by the worker.
1631#[derive(Clone, Copy, Debug)]
1632pub struct WorkerRetryPolicy {
1633    /// Number of retries after the initial request fails.
1634    pub max_retries: usize,
1635    /// Delay before the first retry.
1636    pub initial_backoff: Duration,
1637    /// Maximum delay between retries.
1638    pub max_backoff: Duration,
1639}
1640
1641impl Default for WorkerRetryPolicy {
1642    fn default() -> Self {
1643        Self {
1644            max_retries: 5,
1645            initial_backoff: Duration::from_millis(100),
1646            max_backoff: Duration::from_secs(5),
1647        }
1648    }
1649}
1650
1651#[derive(Clone)]
1652pub struct Worker {
1653    client: Client,
1654    worker_id: String,
1655    task_queue: String,
1656    workflows: HashMap<String, RegisteredWorkflow>,
1657    activities: HashMap<String, ActivityHandler>,
1658    queries: HashMap<String, HashMap<String, RegisteredQuery>>,
1659    max_concurrent_workflow_tasks: usize,
1660    max_concurrent_activity_tasks: usize,
1661    poll_timeout: Duration,
1662    heartbeat_interval: Duration,
1663    retry_policy: WorkerRetryPolicy,
1664    heartbeat_observer: Option<WorkerHeartbeatObserver>,
1665}
1666
1667impl Worker {
1668    pub fn new(client: Client, task_queue: impl Into<String>) -> Self {
1669        Self {
1670            client,
1671            worker_id: default_worker_id(),
1672            task_queue: task_queue.into(),
1673            workflows: HashMap::new(),
1674            activities: HashMap::new(),
1675            queries: HashMap::new(),
1676            max_concurrent_workflow_tasks: 10,
1677            max_concurrent_activity_tasks: 10,
1678            poll_timeout: Duration::from_secs(30),
1679            heartbeat_interval: Duration::from_secs(60),
1680            retry_policy: WorkerRetryPolicy::default(),
1681            heartbeat_observer: None,
1682        }
1683    }
1684
1685    pub fn worker_id(mut self, worker_id: impl Into<String>) -> Self {
1686        self.worker_id = worker_id.into();
1687        self
1688    }
1689
1690    pub fn poll_timeout(mut self, timeout: Duration) -> Self {
1691        self.poll_timeout = timeout;
1692        self
1693    }
1694
1695    pub fn heartbeat_interval(mut self, interval: Duration) -> Self {
1696        self.heartbeat_interval = interval;
1697        self
1698    }
1699
1700    /// Configure bounded retries for task-poll acquisition and worker heartbeats.
1701    pub fn retry_policy(mut self, policy: WorkerRetryPolicy) -> Self {
1702        self.retry_policy = policy;
1703        self
1704    }
1705
1706    pub fn on_worker_heartbeat<F>(mut self, observer: F) -> Self
1707    where
1708        F: Fn(&WorkerHeartbeatObservation) + Send + Sync + 'static,
1709    {
1710        self.heartbeat_observer = Some(Arc::new(observer));
1711        self
1712    }
1713
1714    pub fn max_concurrent_workflow_tasks(mut self, count: usize) -> Self {
1715        self.max_concurrent_workflow_tasks = count.max(1);
1716        self
1717    }
1718
1719    pub fn max_concurrent_activity_tasks(mut self, count: usize) -> Self {
1720        self.max_concurrent_activity_tasks = count.max(1);
1721        self
1722    }
1723
1724    pub fn register_workflow<F, Fut>(&mut self, workflow_type: impl Into<String>, handler: F)
1725    where
1726        F: Fn(WorkflowContext, Value) -> Fut + Send + Sync + 'static,
1727        Fut: Future<Output = Result<Value>> + Send + 'static,
1728    {
1729        self.workflows.insert(
1730            workflow_type.into(),
1731            RegisteredWorkflow {
1732                execute: Arc::new(move |ctx, input| Box::pin(handler(ctx, input))),
1733                replay: None,
1734                state_type: None,
1735            },
1736        );
1737    }
1738
1739    /// Register a workflow whose typed instance state can be reconstructed for queries.
1740    ///
1741    /// `state_factory` creates a fresh instance for every normal workflow task and
1742    /// query replay. The workflow handler is the single source of truth for state
1743    /// transitions: it updates [`WorkflowInstance`] after activities and signals
1744    /// resolve. Query replay runs this same handler over committed history and
1745    /// discards any commands it would emit.
1746    pub fn register_replayed_workflow<S, Factory, F, Fut>(
1747        &mut self,
1748        workflow_type: impl Into<String>,
1749        state_factory: Factory,
1750        handler: F,
1751    ) where
1752        S: Clone + Send + Sync + 'static,
1753        Factory: Fn() -> S + Send + Sync + 'static,
1754        F: Fn(WorkflowContext, Value, WorkflowInstance<S>) -> Fut + Send + Sync + 'static,
1755        Fut: Future<Output = Result<Value>> + Send + 'static,
1756    {
1757        let state_factory = Arc::new(state_factory);
1758        let handler = Arc::new(handler);
1759
1760        let execute_factory = Arc::clone(&state_factory);
1761        let execute_handler = Arc::clone(&handler);
1762        let execute = Arc::new(move |ctx: WorkflowContext, input: Value| {
1763            let state = WorkflowInstance::new(execute_factory());
1764            let future = execute_handler(ctx, input, state);
1765            Box::pin(future) as WorkflowFuture
1766        });
1767
1768        let replay = Arc::new(move |ctx: WorkflowContext, input: Value| {
1769            let state = WorkflowInstance::new(state_factory());
1770            let snapshot_state = state.clone();
1771            let snapshot: WorkflowStateSnapshot =
1772                Arc::new(move || Ok(Arc::new(snapshot_state.snapshot()?) as ErasedWorkflowState));
1773            let future = handler(ctx, input, state);
1774            ReplayedWorkflowInvocation {
1775                future: Box::pin(future),
1776                snapshot,
1777            }
1778        });
1779
1780        self.workflows.insert(
1781            workflow_type.into(),
1782            RegisteredWorkflow {
1783                execute,
1784                replay: Some(replay),
1785                state_type: Some(TypeId::of::<S>()),
1786            },
1787        );
1788    }
1789
1790    pub fn register_activity<F, Fut>(&mut self, activity_type: impl Into<String>, handler: F)
1791    where
1792        F: Fn(ActivityContext, Value) -> Fut + Send + Sync + 'static,
1793        Fut: Future<Output = Result<Value>> + Send + 'static,
1794    {
1795        self.activities.insert(
1796            activity_type.into(),
1797            Arc::new(move |ctx, args| Box::pin(handler(ctx, args))),
1798        );
1799    }
1800
1801    /// Register a named, read-only query handler for a workflow type.
1802    ///
1803    /// The workflow type must also be registered with [`Worker::register_workflow`]
1804    /// before the worker runs. The handler receives only an immutable committed
1805    /// state snapshot and normalized query arguments.
1806    pub fn register_query<F, Fut>(
1807        &mut self,
1808        workflow_type: impl Into<String>,
1809        query_name: impl Into<String>,
1810        handler: F,
1811    ) where
1812        F: Fn(QueryContext, Value) -> Fut + Send + Sync + 'static,
1813        Fut: Future<Output = Result<Value>> + Send + 'static,
1814    {
1815        self.queries
1816            .entry(workflow_type.into())
1817            .or_default()
1818            .insert(
1819                query_name.into(),
1820                RegisteredQuery::Snapshot(Arc::new(move |ctx, args| Box::pin(handler(ctx, args)))),
1821            );
1822    }
1823
1824    /// Register a named query against deterministically replayed instance state.
1825    ///
1826    /// The workflow type must use [`Worker::register_replayed_workflow`] with the
1827    /// same state type `S`. The handler receives an immutable, detached state
1828    /// clone, so successful and failed queries cannot affect workflow execution
1829    /// or the state reconstructed by a later query.
1830    pub fn register_replayed_query<S, F, Fut>(
1831        &mut self,
1832        workflow_type: impl Into<String>,
1833        query_name: impl Into<String>,
1834        handler: F,
1835    ) where
1836        S: Clone + Send + Sync + 'static,
1837        F: Fn(QueryContext, Arc<S>, Value) -> Fut + Send + Sync + 'static,
1838        Fut: Future<Output = Result<Value>> + Send + 'static,
1839    {
1840        let handler = Arc::new(handler);
1841        let erased_handler: ReplayedQueryHandler = Arc::new(move |ctx, state, args| {
1842            let state = state.downcast::<S>().map_err(|_| {
1843                "registered query state type does not match the replayed workflow state".to_string()
1844            })?;
1845            Ok(Box::pin(handler(ctx, state, args)))
1846        });
1847
1848        self.queries
1849            .entry(workflow_type.into())
1850            .or_default()
1851            .insert(
1852                query_name.into(),
1853                RegisteredQuery::Replayed {
1854                    state_type: TypeId::of::<S>(),
1855                    handler: erased_handler,
1856                },
1857            );
1858    }
1859
1860    pub async fn register(&self) -> Result<RegisterWorkerResponse> {
1861        self.client
1862            .register_worker_with_capabilities(
1863                &self.worker_id,
1864                &self.task_queue,
1865                self.workflows.keys().cloned().collect(),
1866                self.activities.keys().cloned().collect(),
1867                self.max_concurrent_workflow_tasks,
1868                self.max_concurrent_activity_tasks,
1869                (!self.queries.is_empty())
1870                    .then(|| QUERY_TASKS_CAPABILITY.to_string())
1871                    .into_iter()
1872                    .collect(),
1873            )
1874            .await
1875    }
1876
1877    /// Run until shutdown or a terminal worker error occurs.
1878    ///
1879    /// Empty long-poll expirations do not stop the worker. Retryable poll and
1880    /// heartbeat failures use [`WorkerRetryPolicy`] independently, while
1881    /// authentication, protocol, and other non-retryable failures are returned.
1882    pub async fn run(&self) -> Result<()> {
1883        self.run_until(std::future::pending::<()>()).await
1884    }
1885
1886    /// Run until `shutdown` resolves or a terminal worker error occurs.
1887    ///
1888    /// This has the same liveness and terminal-error contract as [`Worker::run`].
1889    pub async fn run_until<F>(&self, shutdown: F) -> Result<()>
1890    where
1891        F: Future<Output = ()>,
1892    {
1893        let registration = self.register().await?;
1894        let mut heartbeat = tokio::time::interval(Duration::from_secs(
1895            registration
1896                .heartbeat_interval_seconds
1897                .unwrap_or(self.heartbeat_interval.as_secs().max(1)),
1898        ));
1899        tokio::pin!(shutdown);
1900        let stop = Arc::new(AtomicBool::new(false));
1901        // Poll responses may already have leased server-side work by the time
1902        // they become ready, so each poller owns its responses through
1903        // completion or failure instead of racing raw polls in this select.
1904        let mut workflow_poller = (!self.workflows.is_empty()).then(|| {
1905            let worker = self.clone();
1906            let stop = Arc::clone(&stop);
1907            tokio::spawn(async move { worker.poll_workflows_until_stopped(stop).await })
1908        });
1909        let mut activity_poller = (!self.activities.is_empty()).then(|| {
1910            let worker = self.clone();
1911            let stop = Arc::clone(&stop);
1912            tokio::spawn(async move { worker.poll_activities_until_stopped(stop).await })
1913        });
1914        let mut query_poller = (!self.queries.is_empty()).then(|| {
1915            let worker = self.clone();
1916            let stop = Arc::clone(&stop);
1917            tokio::spawn(async move { worker.poll_queries_until_stopped(stop).await })
1918        });
1919
1920        loop {
1921            tokio::select! {
1922                _ = &mut shutdown => {
1923                    stop.store(true, Ordering::SeqCst);
1924                    break;
1925                }
1926                _ = heartbeat.tick() => {
1927                    match self.retry_worker_operation(|| {
1928                        self.client.heartbeat_worker(
1929                            &self.worker_id,
1930                            self.max_concurrent_workflow_tasks,
1931                            self.max_concurrent_activity_tasks,
1932                        )
1933                    }).await
1934                    {
1935                        Ok(acknowledgement) => {
1936                            if let Some(observer) = &self.heartbeat_observer {
1937                                observer(&WorkerHeartbeatObservation {
1938                                    worker_id: self.worker_id.clone(),
1939                                    task_queue: self.task_queue.clone(),
1940                                    acknowledged_at_unix_millis: SystemTime::now()
1941                                        .duration_since(UNIX_EPOCH)
1942                                        .unwrap_or_default()
1943                                        .as_millis()
1944                                        .min(u64::MAX as u128)
1945                                        as u64,
1946                                    acknowledgement,
1947                                });
1948                            }
1949                        }
1950                        Err(error) => {
1951                            stop.store(true, Ordering::SeqCst);
1952                            join_pollers(workflow_poller.take(), activity_poller.take(), query_poller.take()).await?;
1953                            return Err(error);
1954                        }
1955                    }
1956                }
1957                result = OptionFuture::from(workflow_poller.as_mut()), if workflow_poller.is_some() => {
1958                    workflow_poller = None;
1959                    stop.store(true, Ordering::SeqCst);
1960                    let poller_result = optional_poller_result("workflow", result);
1961                    let join_result =
1962                        join_pollers(workflow_poller.take(), activity_poller.take(), query_poller.take()).await;
1963                    poller_result?;
1964                    join_result?;
1965                    return Err(Error::WorkerLoop(
1966                        "workflow poller stopped unexpectedly".to_string(),
1967                    ));
1968                }
1969                result = OptionFuture::from(activity_poller.as_mut()), if activity_poller.is_some() => {
1970                    activity_poller = None;
1971                    stop.store(true, Ordering::SeqCst);
1972                    let poller_result = optional_poller_result("activity", result);
1973                    let join_result =
1974                        join_pollers(workflow_poller.take(), activity_poller.take(), query_poller.take()).await;
1975                    poller_result?;
1976                    join_result?;
1977                    return Err(Error::WorkerLoop(
1978                        "activity poller stopped unexpectedly".to_string(),
1979                    ));
1980                }
1981                result = OptionFuture::from(query_poller.as_mut()), if query_poller.is_some() => {
1982                    query_poller = None;
1983                    stop.store(true, Ordering::SeqCst);
1984                    let poller_result = optional_poller_result("query", result);
1985                    let join_result =
1986                        join_pollers(workflow_poller.take(), activity_poller.take(), query_poller.take()).await;
1987                    poller_result?;
1988                    join_result?;
1989                    return Err(Error::WorkerLoop(
1990                        "query poller stopped unexpectedly".to_string(),
1991                    ));
1992                }
1993            }
1994        }
1995
1996        join_pollers(
1997            workflow_poller.take(),
1998            activity_poller.take(),
1999            query_poller.take(),
2000        )
2001        .await
2002    }
2003
2004    pub async fn run_once(&self) -> Result<usize> {
2005        let mut handled = 0;
2006        if self.poll_workflow_once().await? {
2007            handled += 1;
2008        }
2009        if self.poll_activity_once().await? {
2010            handled += 1;
2011        }
2012        if !self.queries.is_empty() && self.poll_query_once().await? {
2013            handled += 1;
2014        }
2015        Ok(handled)
2016    }
2017
2018    async fn poll_workflow_once(&self) -> Result<bool> {
2019        let Some(task) = self
2020            .retry_worker_operation(|| {
2021                self.client
2022                    .poll_workflow_task(&self.worker_id, &self.task_queue, self.poll_timeout)
2023            })
2024            .await?
2025        else {
2026            return Ok(false);
2027        };
2028
2029        let task_id = task.task_id.clone();
2030        let attempt = task.workflow_task_attempt;
2031        let lease_owner = task
2032            .lease_owner
2033            .clone()
2034            .unwrap_or_else(|| self.worker_id.clone());
2035
2036        match self.execute_workflow_task(task) {
2037            Ok(commands) => {
2038                self.client
2039                    .complete_workflow_task(&task_id, &lease_owner, attempt, commands)
2040                    .await?;
2041            }
2042            Err(error) => {
2043                self.client
2044                    .fail_workflow_task(&task_id, &lease_owner, attempt, error.to_string())
2045                    .await?;
2046            }
2047        }
2048
2049        Ok(true)
2050    }
2051
2052    async fn poll_workflows_until_stopped(self, stop: Arc<AtomicBool>) -> Result<()> {
2053        while !stop.load(Ordering::SeqCst) {
2054            self.poll_workflow_once().await?;
2055        }
2056
2057        Ok(())
2058    }
2059
2060    async fn poll_activity_once(&self) -> Result<bool> {
2061        let Some(task) = self
2062            .retry_worker_operation(|| {
2063                self.client
2064                    .poll_activity_task(&self.worker_id, &self.task_queue, self.poll_timeout)
2065            })
2066            .await?
2067        else {
2068            return Ok(false);
2069        };
2070
2071        let task_id = task.task_id.clone();
2072        let attempt_id = task
2073            .activity_attempt_id
2074            .clone()
2075            .or(task.attempt_id.clone())
2076            .unwrap_or_default();
2077        let lease_owner = task
2078            .lease_owner
2079            .clone()
2080            .unwrap_or_else(|| self.worker_id.clone());
2081        let codec = task.payload_codec.clone();
2082        let result = self.execute_activity_task(task).await;
2083        match result {
2084            Ok(value) => {
2085                self.client
2086                    .complete_activity_task(&task_id, &attempt_id, &lease_owner, value, &codec)
2087                    .await?;
2088            }
2089            Err(error) => {
2090                self.client
2091                    .fail_activity_task(
2092                        &task_id,
2093                        &attempt_id,
2094                        &lease_owner,
2095                        error.to_string(),
2096                        false,
2097                    )
2098                    .await?;
2099            }
2100        }
2101
2102        Ok(true)
2103    }
2104
2105    async fn poll_activities_until_stopped(self, stop: Arc<AtomicBool>) -> Result<()> {
2106        while !stop.load(Ordering::SeqCst) {
2107            self.poll_activity_once().await?;
2108        }
2109
2110        Ok(())
2111    }
2112
2113    async fn poll_query_once(&self) -> Result<bool> {
2114        let Some(task) = self
2115            .retry_worker_operation(|| {
2116                self.client
2117                    .poll_query_task(&self.worker_id, &self.task_queue, self.poll_timeout)
2118            })
2119            .await?
2120        else {
2121            return Ok(false);
2122        };
2123
2124        let query_task_id = task.query_task_id.clone();
2125        let attempt = task.query_task_attempt;
2126        let lease_owner = task
2127            .lease_owner
2128            .clone()
2129            .unwrap_or_else(|| self.worker_id.clone());
2130        let codec = task.payload_codec.clone();
2131
2132        match self.execute_query_task(task).await {
2133            Ok(value) => {
2134                let result_envelope = match encode_value_envelope(&value, &codec) {
2135                    Ok(result_envelope) => result_envelope,
2136                    Err(error) => {
2137                        let failure = self
2138                            .client
2139                            .fail_query_task(
2140                                &query_task_id,
2141                                &lease_owner,
2142                                attempt,
2143                                error.to_string(),
2144                                "query_result_encode_failed",
2145                                "QueryResultEncodeFailed",
2146                            )
2147                            .await;
2148                        if let Err(error) = failure {
2149                            if !query_task_rejection_is_final(&error) {
2150                                return Err(error);
2151                            }
2152                        }
2153                        return Ok(true);
2154                    }
2155                };
2156
2157                if let Err(error) = self
2158                    .client
2159                    .complete_query_task_with_envelope(
2160                        &query_task_id,
2161                        &lease_owner,
2162                        attempt,
2163                        value,
2164                        result_envelope,
2165                    )
2166                    .await
2167                {
2168                    if !query_task_rejection_is_final(&error) {
2169                        return Err(error);
2170                    }
2171                }
2172            }
2173            Err(failure) => {
2174                let result = self
2175                    .client
2176                    .fail_query_task(
2177                        &query_task_id,
2178                        &lease_owner,
2179                        attempt,
2180                        failure.message,
2181                        failure.reason,
2182                        failure.failure_type,
2183                    )
2184                    .await;
2185                if let Err(error) = result {
2186                    if !query_task_rejection_is_final(&error) {
2187                        return Err(error);
2188                    }
2189                }
2190            }
2191        }
2192
2193        Ok(true)
2194    }
2195
2196    async fn poll_queries_until_stopped(self, stop: Arc<AtomicBool>) -> Result<()> {
2197        while !stop.load(Ordering::SeqCst) {
2198            self.poll_query_once().await?;
2199        }
2200
2201        Ok(())
2202    }
2203
2204    async fn retry_worker_operation<T, F, Fut>(&self, mut operation: F) -> Result<T>
2205    where
2206        F: FnMut() -> Fut,
2207        Fut: Future<Output = Result<T>>,
2208    {
2209        let mut retries = 0;
2210
2211        loop {
2212            match operation().await {
2213                Err(error)
2214                    if worker_operation_is_retryable(&error)
2215                        && retries < self.retry_policy.max_retries =>
2216                {
2217                    retries += 1;
2218                    tokio::time::sleep(worker_retry_delay(self.retry_policy, retries)).await;
2219                }
2220                result => return result,
2221            }
2222        }
2223    }
2224
2225    async fn execute_query_task(
2226        &self,
2227        mut task: QueryTask,
2228    ) -> std::result::Result<Value, QueryTaskExecutionFailure> {
2229        if !matches!(task.payload_codec.as_str(), DEFAULT_CODEC | JSON_CODEC) {
2230            return Err(QueryTaskExecutionFailure::new(
2231                "query_payload_decode_failed",
2232                format!(
2233                    "cannot decode query payload with unsupported codec {:?}",
2234                    task.payload_codec
2235                ),
2236                "QueryPayloadDecodeFailed",
2237            ));
2238        }
2239
2240        if !self.workflows.contains_key(&task.workflow_type) {
2241            return Err(QueryTaskExecutionFailure::new(
2242                "query_workflow_type_not_registered",
2243                format!("no workflow registered for type {:?}", task.workflow_type),
2244                "WorkflowTypeNotRegistered",
2245            ));
2246        }
2247
2248        let Some(handlers) = self.queries.get(&task.workflow_type) else {
2249            return Err(QueryTaskExecutionFailure::new(
2250                "query_handler_unavailable",
2251                format!(
2252                    "query handlers are unavailable for workflow type {:?}",
2253                    task.workflow_type
2254                ),
2255                "QueryHandlerUnavailable",
2256            ));
2257        };
2258        let Some(query) = handlers.get(&task.query_name) else {
2259            return Err(QueryTaskExecutionFailure::new(
2260                "rejected_unknown_query",
2261                format!("unknown query {:?}", task.query_name),
2262                "QueryFailed",
2263            ));
2264        };
2265
2266        let args = decode_task_arguments(task.query_arguments.as_ref(), &task.payload_codec)
2267            .map_err(|error| {
2268                QueryTaskExecutionFailure::new(
2269                    "query_payload_decode_failed",
2270                    format!("cannot decode query arguments: {error}"),
2271                    "QueryPayloadDecodeFailed",
2272                )
2273            })?;
2274        let workflow_input =
2275            decode_task_arguments(task.workflow_arguments.as_ref(), &task.payload_codec).map_err(
2276                |error| {
2277                    QueryTaskExecutionFailure::new(
2278                        "query_workflow_state_unavailable",
2279                        format!("cannot decode workflow start input: {error}"),
2280                        "QueryWorkflowStateUnavailable",
2281                    )
2282                },
2283            )?;
2284        hydrate_query_history_from_export(&mut task).map_err(|error| {
2285            QueryTaskExecutionFailure::new(
2286                "query_workflow_state_unavailable",
2287                format!("cannot restore query history snapshot: {error}"),
2288                "QueryWorkflowStateUnavailable",
2289            )
2290        })?;
2291        enrich_query_history_from_export(&mut task).map_err(|error| {
2292            QueryTaskExecutionFailure::new(
2293                "query_workflow_state_unavailable",
2294                format!("cannot restore compact query history payloads: {error}"),
2295                "QueryWorkflowStateUnavailable",
2296            )
2297        })?;
2298        let signal_events = query_signal_events(&task).map_err(|error| {
2299            QueryTaskExecutionFailure::new(
2300                "query_workflow_state_unavailable",
2301                format!("cannot decode committed workflow signals: {error}"),
2302                "QueryWorkflowStateUnavailable",
2303            )
2304        })?;
2305        let history_events = Arc::new(std::mem::take(&mut task.history_events));
2306        let context = QueryContext {
2307            workflow_id: task.workflow_id,
2308            run_id: task.run_id,
2309            workflow_type: task.workflow_type.clone(),
2310            run_status: task.run_status,
2311            workflow_input,
2312            history_events: Arc::clone(&history_events),
2313            signal_events: Arc::new(signal_events),
2314        };
2315
2316        let future = match query {
2317            RegisteredQuery::Snapshot(handler) => handler(context, args),
2318            RegisteredQuery::Replayed {
2319                state_type,
2320                handler,
2321            } => {
2322                let workflow = self
2323                    .workflows
2324                    .get(&task.workflow_type)
2325                    .expect("workflow registration was checked above");
2326                if workflow.state_type != Some(*state_type) {
2327                    return Err(QueryTaskExecutionFailure::new(
2328                        "query_workflow_state_unavailable",
2329                        "replayed query state type does not match its workflow registration",
2330                        "QueryWorkflowStateUnavailable",
2331                    ));
2332                }
2333                let replay = workflow.replay.as_ref().ok_or_else(|| {
2334                    QueryTaskExecutionFailure::new(
2335                        "query_workflow_state_unavailable",
2336                        format!(
2337                            "workflow type {:?} is not registered for instance-state replay",
2338                            task.workflow_type
2339                        ),
2340                        "QueryWorkflowStateUnavailable",
2341                    )
2342                })?;
2343                let workflow_state = Arc::new(Mutex::new(
2344                    WorkflowState::new_with_identity(
2345                        history_events.as_ref().clone(),
2346                        context.workflow_id.clone(),
2347                        context.run_id.clone(),
2348                        self.task_queue.clone(),
2349                        task.payload_codec,
2350                        None,
2351                    )
2352                    .map_err(|error| {
2353                        QueryTaskExecutionFailure::new(
2354                            "query_workflow_state_unavailable",
2355                            format!("workflow replay failed before query: {error}"),
2356                            "QueryWorkflowStateUnavailable",
2357                        )
2358                    })?,
2359                ));
2360                let workflow_context = WorkflowContext {
2361                    state: workflow_state,
2362                };
2363                let mut invocation =
2364                    replay(workflow_context.clone(), context.workflow_input.clone());
2365                let mut cx = TaskContext::from_waker(noop_waker_ref());
2366                match invocation.future.as_mut().poll(&mut cx) {
2367                    Poll::Ready(Ok(_)) => {
2368                        workflow_context
2369                            .ensure_history_consumed()
2370                            .map_err(|error| {
2371                                QueryTaskExecutionFailure::new(
2372                                    "query_workflow_state_unavailable",
2373                                    format!("workflow replay failed before query: {error}"),
2374                                    "QueryWorkflowStateUnavailable",
2375                                )
2376                            })?;
2377                    }
2378                    Poll::Ready(Err(error)) => {
2379                        return Err(QueryTaskExecutionFailure::new(
2380                            "query_workflow_state_unavailable",
2381                            format!("workflow replay failed before query: {error}"),
2382                            "QueryWorkflowStateUnavailable",
2383                        ));
2384                    }
2385                    Poll::Pending => {
2386                        let commands = workflow_context.take_commands().map_err(|error| {
2387                            QueryTaskExecutionFailure::new(
2388                                "query_workflow_state_unavailable",
2389                                format!("workflow replay failed before query: {error}"),
2390                                "QueryWorkflowStateUnavailable",
2391                            )
2392                        })?;
2393                        if commands.is_empty()
2394                            && !workflow_context
2395                                .matched_recorded_pending()
2396                                .map_err(|error| {
2397                                    QueryTaskExecutionFailure::new(
2398                                        "query_workflow_state_unavailable",
2399                                        format!("workflow replay failed before query: {error}"),
2400                                        "QueryWorkflowStateUnavailable",
2401                                    )
2402                                })?
2403                        {
2404                            return Err(QueryTaskExecutionFailure::new(
2405                                "query_workflow_state_unavailable",
2406                                "workflow replay yielded without a durable command",
2407                                "QueryWorkflowStateUnavailable",
2408                            ));
2409                        }
2410                    }
2411                }
2412                let state = (invocation.snapshot)().map_err(|error| {
2413                    QueryTaskExecutionFailure::new(
2414                        "query_workflow_state_unavailable",
2415                        format!("cannot snapshot replayed workflow state: {error}"),
2416                        "QueryWorkflowStateUnavailable",
2417                    )
2418                })?;
2419                handler(context, state, args).map_err(|message| {
2420                    QueryTaskExecutionFailure::new(
2421                        "query_workflow_state_unavailable",
2422                        message,
2423                        "QueryWorkflowStateUnavailable",
2424                    )
2425                })?
2426            }
2427        };
2428
2429        future.await.map_err(|error| {
2430            QueryTaskExecutionFailure::new("query_rejected", error.to_string(), "QueryFailed")
2431        })
2432    }
2433
2434    fn execute_workflow_task(&self, task: WorkflowTask) -> Result<Vec<Value>> {
2435        let workflow = self
2436            .workflows
2437            .get(&task.workflow_type)
2438            .ok_or_else(|| Error::WorkflowNotRegistered(task.workflow_type.clone()))?;
2439        let input = decode_task_arguments(task.arguments.as_ref(), &task.payload_codec)?;
2440        let resume_signal = decode_resume_signal(&task)?;
2441        let state = Arc::new(Mutex::new(WorkflowState::new_with_identity(
2442            task.history_events,
2443            task.workflow_id,
2444            task.run_id,
2445            self.task_queue.clone(),
2446            task.payload_codec.clone(),
2447            resume_signal,
2448        )?));
2449        let ctx = WorkflowContext { state };
2450        let mut future = (workflow.execute)(ctx.clone(), input);
2451        let mut cx = TaskContext::from_waker(noop_waker_ref());
2452
2453        match future.as_mut().poll(&mut cx) {
2454            Poll::Ready(Ok(result)) => {
2455                ctx.ensure_history_consumed()?;
2456                let result = encode_value_envelope(&result, &task.payload_codec)?;
2457                Ok(vec![json!({
2458                    "type": "complete_workflow",
2459                    "result": result
2460                })])
2461            }
2462            Poll::Ready(Err(error @ Error::ChildWorkflowFailed(_))) => {
2463                Ok(vec![workflow_failure_command(&error)])
2464            }
2465            Poll::Ready(Err(error)) => Err(error),
2466            Poll::Pending => {
2467                let commands = ctx.take_commands()?;
2468                if commands.is_empty() && !ctx.matched_recorded_pending()? {
2469                    Err(Error::WorkflowYieldedWithoutCommand)
2470                } else {
2471                    Ok(commands)
2472                }
2473            }
2474        }
2475    }
2476
2477    async fn execute_activity_task(&self, task: ActivityTask) -> Result<Value> {
2478        let handler = self
2479            .activities
2480            .get(&task.activity_type)
2481            .ok_or_else(|| Error::ActivityNotRegistered(task.activity_type.clone()))?;
2482        let args = decode_task_arguments(task.arguments.as_ref(), &task.payload_codec)?;
2483        let attempt_id = task
2484            .activity_attempt_id
2485            .clone()
2486            .or(task.attempt_id.clone())
2487            .unwrap_or_default();
2488        let lease_owner = task
2489            .lease_owner
2490            .clone()
2491            .unwrap_or_else(|| self.worker_id.clone());
2492        let ctx = ActivityContext {
2493            client: self.client.clone(),
2494            task_id: task.task_id,
2495            activity_attempt_id: attempt_id,
2496            lease_owner,
2497            activity_type: task.activity_type,
2498            attempt_number: task.attempt_number,
2499            task_queue: self.task_queue.clone(),
2500            worker_id: self.worker_id.clone(),
2501        };
2502
2503        handler(ctx, args).await
2504    }
2505}
2506
2507fn poller_result(
2508    kind: &str,
2509    result: std::result::Result<Result<()>, tokio::task::JoinError>,
2510) -> Result<()> {
2511    match result {
2512        Ok(result) => result,
2513        Err(error) => Err(Error::WorkerLoop(format!(
2514            "{kind} poller join error: {error}"
2515        ))),
2516    }
2517}
2518
2519fn optional_poller_result(
2520    kind: &str,
2521    result: Option<std::result::Result<Result<()>, tokio::task::JoinError>>,
2522) -> Result<()> {
2523    match result {
2524        Some(result) => poller_result(kind, result),
2525        None => Ok(()),
2526    }
2527}
2528
2529async fn join_pollers(
2530    workflow_poller: Option<tokio::task::JoinHandle<Result<()>>>,
2531    activity_poller: Option<tokio::task::JoinHandle<Result<()>>>,
2532    query_poller: Option<tokio::task::JoinHandle<Result<()>>>,
2533) -> Result<()> {
2534    let mut first_error = None;
2535
2536    if let Some(handle) = workflow_poller {
2537        if let Err(error) = poller_result("workflow", handle.await) {
2538            first_error.get_or_insert(error);
2539        }
2540    }
2541
2542    if let Some(handle) = activity_poller {
2543        if let Err(error) = poller_result("activity", handle.await) {
2544            first_error.get_or_insert(error);
2545        }
2546    }
2547
2548    if let Some(handle) = query_poller {
2549        if let Err(error) = poller_result("query", handle.await) {
2550            first_error.get_or_insert(error);
2551        }
2552    }
2553
2554    if let Some(error) = first_error {
2555        Err(error)
2556    } else {
2557        Ok(())
2558    }
2559}
2560
2561fn default_worker_id() -> String {
2562    let millis = SystemTime::now()
2563        .duration_since(UNIX_EPOCH)
2564        .unwrap_or_default()
2565        .as_millis();
2566    format!("rust-worker-{}-{millis}", std::process::id())
2567}
2568
2569fn unique_request_id(prefix: &str) -> String {
2570    let nanos = SystemTime::now()
2571        .duration_since(UNIX_EPOCH)
2572        .unwrap_or_default()
2573        .as_nanos();
2574    format!("{prefix}-{}-{nanos}", std::process::id())
2575}
2576
2577#[derive(Debug)]
2578struct QueryTaskExecutionFailure {
2579    reason: String,
2580    message: String,
2581    failure_type: String,
2582}
2583
2584impl QueryTaskExecutionFailure {
2585    fn new(
2586        reason: impl Into<String>,
2587        message: impl Into<String>,
2588        failure_type: impl Into<String>,
2589    ) -> Self {
2590        Self {
2591            reason: reason.into(),
2592            message: message.into(),
2593            failure_type: failure_type.into(),
2594        }
2595    }
2596}
2597
2598/// Typed local state owned by one deterministic workflow invocation.
2599///
2600/// Use [`WorkflowInstance::update`] for the same state transitions during
2601/// ordinary execution and replay. A replayed query receives a detached
2602/// immutable `Arc<S>` rather than this mutation-capable handle.
2603#[derive(Clone, Debug)]
2604pub struct WorkflowInstance<S> {
2605    state: Arc<Mutex<S>>,
2606}
2607
2608impl<S> WorkflowInstance<S> {
2609    fn new(state: S) -> Self {
2610        Self {
2611            state: Arc::new(Mutex::new(state)),
2612        }
2613    }
2614
2615    /// Read the current workflow-instance state without changing it.
2616    pub fn read<R>(&self, reader: impl FnOnce(&S) -> R) -> Result<R> {
2617        let state = self
2618            .state
2619            .lock()
2620            .map_err(|_| Error::WorkflowStatePoisoned)?;
2621        Ok(reader(&state))
2622    }
2623
2624    /// Apply one deterministic workflow-instance state transition.
2625    pub fn update<R>(&self, transition: impl FnOnce(&mut S) -> R) -> Result<R> {
2626        let mut state = self
2627            .state
2628            .lock()
2629            .map_err(|_| Error::WorkflowStatePoisoned)?;
2630        Ok(transition(&mut state))
2631    }
2632}
2633
2634impl<S: Clone> WorkflowInstance<S> {
2635    fn snapshot(&self) -> Result<S> {
2636        self.read(Clone::clone)
2637    }
2638}
2639
2640#[derive(Clone, Debug)]
2641pub struct WorkflowContext {
2642    state: Arc<Mutex<WorkflowState>>,
2643}
2644
2645impl WorkflowContext {
2646    /// Identity of the parent workflow currently being replayed.
2647    pub fn workflow_identity(&self) -> Result<WorkflowIdentity> {
2648        let state = self
2649            .state
2650            .lock()
2651            .map_err(|_| Error::WorkflowStatePoisoned)?;
2652        Ok(WorkflowIdentity {
2653            workflow_id: state.workflow_id.clone(),
2654            run_id: state.run_id.clone(),
2655        })
2656    }
2657
2658    pub fn activity<T: Serialize>(
2659        &self,
2660        activity_type: impl Into<String>,
2661        args: T,
2662    ) -> ActivityCall {
2663        self.activity_on_queue(activity_type, None::<String>, args)
2664    }
2665
2666    pub fn activity_on_queue<T, Q>(
2667        &self,
2668        activity_type: impl Into<String>,
2669        task_queue: Option<Q>,
2670        args: T,
2671    ) -> ActivityCall
2672    where
2673        T: Serialize,
2674        Q: Into<String>,
2675    {
2676        ActivityCall {
2677            ctx: self.clone(),
2678            activity_type: activity_type.into(),
2679            task_queue: task_queue.map(Into::into),
2680            args: Some(serde_json::to_value(args).map_err(Error::from)),
2681            scheduled: false,
2682        }
2683    }
2684
2685    pub fn wait_signal(&self, signal_name: impl Into<String>) -> SignalCall {
2686        SignalCall {
2687            ctx: self.clone(),
2688            signal_name: signal_name.into(),
2689            opened_wait: false,
2690            matched_pending: false,
2691        }
2692    }
2693
2694    /// Wait for server-backed durable time without blocking the worker executor.
2695    ///
2696    /// Polling this future emits one `start_timer` command and yields. The
2697    /// server records the deadline, so neither worker nor server restarts reset
2698    /// the wait. Replay resolves the future only from a `TimerScheduled` and
2699    /// `TimerFired` pair at the same position in the shared durable-command
2700    /// stream, with matching sequence, timer identity, and delay. Sub-second
2701    /// durations round up because protocol deadlines use whole seconds.
2702    ///
2703    /// ```no_run
2704    /// # use durable_workflow::{json, Client, Worker};
2705    /// # use std::time::Duration;
2706    /// # fn configure(client: Client) {
2707    /// let mut worker = Worker::new(client, "rust-workers");
2708    /// worker.register_workflow("delayed-greeting", |ctx, _input| async move {
2709    ///     ctx.sleep(Duration::from_secs(5)).await?;
2710    ///     Ok(json!({"status": "timer fired"}))
2711    /// });
2712    /// # }
2713    /// ```
2714    pub fn sleep(&self, duration: Duration) -> TimerCall {
2715        let delay_seconds = duration
2716            .as_secs()
2717            .checked_add(u64::from(duration.subsec_nanos() > 0));
2718        TimerCall {
2719            ctx: self.clone(),
2720            delay_seconds,
2721            scheduled: false,
2722            matched_pending: false,
2723        }
2724    }
2725
2726    /// Alias for [`WorkflowContext::sleep`] for timer-oriented workflow code.
2727    pub fn start_timer(&self, duration: Duration) -> TimerCall {
2728        self.sleep(duration)
2729    }
2730
2731    /// Start a named durable child on an explicit queue and await its result.
2732    ///
2733    /// The command is recorded in the parent's sequence-ordered durable command
2734    /// stream. Replay keeps a scheduled child pending without emitting another
2735    /// start, or consumes its matching terminal `ChildRun*` outcome. Successful
2736    /// values preserve the history payload codec and include both sides of the
2737    /// durable relationship; failures are returned as
2738    /// [`Error::ChildWorkflowFailed`].
2739    ///
2740    /// ```no_run
2741    /// # use durable_workflow::{json, ChildWorkflowOptions, Client, ParentClosePolicy, Worker};
2742    /// # fn configure(client: Client) {
2743    /// let mut worker = Worker::new(client, "parent-workers");
2744    /// worker.register_workflow("order-parent", |ctx, _input| async move {
2745    ///     let child = ctx
2746    ///         .start_child_workflow(
2747    ///             "fulfil-order",
2748    ///             ChildWorkflowOptions::new("fulfilment-workers")
2749    ///                 .parent_close_policy(ParentClosePolicy::RequestCancel),
2750    ///             json!([{"order_id": "order-42"}]),
2751    ///         )
2752    ///         .await?;
2753    ///     Ok(child.result)
2754    /// });
2755    /// # }
2756    /// ```
2757    pub fn start_child_workflow<T: Serialize>(
2758        &self,
2759        workflow_type: impl Into<String>,
2760        options: ChildWorkflowOptions,
2761        args: T,
2762    ) -> ChildWorkflowCall {
2763        ChildWorkflowCall {
2764            ctx: self.clone(),
2765            workflow_type: workflow_type.into(),
2766            options,
2767            args: Some(serde_json::to_value(args).map_err(Error::from)),
2768            scheduled: false,
2769            matched_pending: false,
2770        }
2771    }
2772
2773    fn take_commands(&self) -> Result<Vec<Value>> {
2774        let mut state = self
2775            .state
2776            .lock()
2777            .map_err(|_| Error::WorkflowStatePoisoned)?;
2778        Ok(std::mem::take(&mut state.commands))
2779    }
2780
2781    fn matched_recorded_pending(&self) -> Result<bool> {
2782        let state = self
2783            .state
2784            .lock()
2785            .map_err(|_| Error::WorkflowStatePoisoned)?;
2786        Ok(state.matched_recorded_pending)
2787    }
2788
2789    fn ensure_history_consumed(&self) -> Result<()> {
2790        let state = self
2791            .state
2792            .lock()
2793            .map_err(|_| Error::WorkflowStatePoisoned)?;
2794        if let Some(command) = state.recorded_commands.get(state.command_cursor) {
2795            return Err(Error::NonDeterministicReplay(ReplayFailure::new(
2796                "recorded_commands_unconsumed",
2797                Some(command.sequence()),
2798                Some(command.shape().to_string()),
2799                Some("workflow completion".to_string()),
2800                "workflow completed before consuming all recorded durable commands",
2801            )));
2802        }
2803        Ok(())
2804    }
2805}
2806
2807#[derive(Debug)]
2808struct WorkflowState {
2809    history: Vec<HistoryEvent>,
2810    workflow_id: Option<String>,
2811    run_id: Option<String>,
2812    task_queue: String,
2813    payload_codec: String,
2814    resume_signal: Option<ResumeSignal>,
2815    recorded_commands: Vec<RecordedCommand>,
2816    command_cursor: usize,
2817    matched_recorded_pending: bool,
2818    signal_cursors: HashMap<String, usize>,
2819    commands: Vec<Value>,
2820}
2821
2822impl WorkflowState {
2823    #[cfg(test)]
2824    fn new(
2825        history: Vec<HistoryEvent>,
2826        task_queue: String,
2827        payload_codec: String,
2828        resume_signal: Option<ResumeSignal>,
2829    ) -> Result<Self> {
2830        Self::new_with_identity(
2831            history,
2832            None,
2833            None,
2834            task_queue,
2835            payload_codec,
2836            resume_signal,
2837        )
2838    }
2839
2840    fn new_with_identity(
2841        history: Vec<HistoryEvent>,
2842        workflow_id: Option<String>,
2843        run_id: Option<String>,
2844        task_queue: String,
2845        payload_codec: String,
2846        resume_signal: Option<ResumeSignal>,
2847    ) -> Result<Self> {
2848        let recorded_commands = recorded_commands(
2849            &history,
2850            &payload_codec,
2851            WorkflowIdentity {
2852                workflow_id: workflow_id.clone(),
2853                run_id: run_id.clone(),
2854            },
2855        )?;
2856        Ok(Self {
2857            history,
2858            workflow_id,
2859            run_id,
2860            task_queue,
2861            payload_codec,
2862            resume_signal,
2863            recorded_commands,
2864            command_cursor: 0,
2865            matched_recorded_pending: false,
2866            signal_cursors: HashMap::new(),
2867            commands: Vec::new(),
2868        })
2869    }
2870}
2871
2872#[derive(Clone, Debug)]
2873enum RecordedCommand {
2874    Activity {
2875        sequence: u64,
2876        activity_type: Option<String>,
2877        result: Option<Value>,
2878    },
2879    Timer {
2880        sequence: u64,
2881        delay_seconds: u64,
2882        fired: bool,
2883    },
2884    ChildWorkflow {
2885        sequence: u64,
2886        workflow_type: Option<String>,
2887        outcome: Option<ChildWorkflowOutcome>,
2888    },
2889    SignalWait {
2890        sequence: u64,
2891        signal_name: Option<String>,
2892    },
2893}
2894
2895impl RecordedCommand {
2896    fn sequence(&self) -> u64 {
2897        match self {
2898            Self::Activity { sequence, .. }
2899            | Self::Timer { sequence, .. }
2900            | Self::ChildWorkflow { sequence, .. }
2901            | Self::SignalWait { sequence, .. } => *sequence,
2902        }
2903    }
2904
2905    fn shape(&self) -> &'static str {
2906        match self {
2907            Self::Activity { .. } => "activity",
2908            Self::Timer { .. } => "timer",
2909            Self::ChildWorkflow { .. } => "child workflow",
2910            Self::SignalWait { .. } => "signal wait",
2911        }
2912    }
2913}
2914
2915#[derive(Clone, Debug)]
2916struct ResumeSignal {
2917    signal_id: Option<String>,
2918    signal_name: String,
2919    arguments: Vec<Value>,
2920}
2921
2922pub struct ActivityCall {
2923    ctx: WorkflowContext,
2924    activity_type: String,
2925    task_queue: Option<String>,
2926    args: Option<Result<Value>>,
2927    scheduled: bool,
2928}
2929
2930impl Future for ActivityCall {
2931    type Output = Result<Value>;
2932
2933    fn poll(mut self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
2934        let ctx = self.ctx.clone();
2935        let mut state = match ctx.state.lock() {
2936            Ok(state) => state,
2937            Err(_) => return Poll::Ready(Err(Error::WorkflowStatePoisoned)),
2938        };
2939
2940        if self.scheduled {
2941            return Poll::Pending;
2942        }
2943
2944        if let Some(recorded) = state.recorded_commands.get(state.command_cursor).cloned() {
2945            let sequence = recorded.sequence();
2946            match recorded {
2947                RecordedCommand::Activity {
2948                    activity_type,
2949                    result,
2950                    ..
2951                } => {
2952                    if let Some(recorded_type) = activity_type {
2953                        if recorded_type != self.activity_type {
2954                            return Poll::Ready(Err(Error::NonDeterministicReplay(
2955                                ReplayFailure::new(
2956                                    "recorded_command_detail_mismatch",
2957                                    Some(sequence),
2958                                    Some(format!("activity:{recorded_type}")),
2959                                    Some(format!("activity:{}", self.activity_type)),
2960                                    "recorded activity type differs from the current workflow command",
2961                                ),
2962                            )));
2963                        }
2964                    }
2965                    state.command_cursor += 1;
2966                    if let Some(value) = result {
2967                        return Poll::Ready(Ok(value));
2968                    }
2969                    state.matched_recorded_pending = true;
2970                    self.scheduled = true;
2971                    return Poll::Pending;
2972                }
2973                other => {
2974                    return Poll::Ready(Err(command_mismatch(
2975                        &other,
2976                        format!("activity:{}", self.activity_type),
2977                    )));
2978                }
2979            }
2980        }
2981
2982        if !self.scheduled {
2983            let task_queue = self
2984                .task_queue
2985                .clone()
2986                .unwrap_or_else(|| state.task_queue.clone());
2987            let args = match self.args.take().unwrap_or(Ok(Value::Null)) {
2988                Ok(args) => args,
2989                Err(error) => return Poll::Ready(Err(error)),
2990            };
2991            let arguments = normalize_arguments(args);
2992            let envelope = match encode_value_envelope(&arguments, &state.payload_codec) {
2993                Ok(envelope) => envelope,
2994                Err(error) => return Poll::Ready(Err(error)),
2995            };
2996
2997            state.commands.push(json!({
2998                "type": "schedule_activity",
2999                "activity_type": self.activity_type.clone(),
3000                "queue": task_queue,
3001                "arguments": envelope
3002            }));
3003            self.scheduled = true;
3004        }
3005
3006        Poll::Pending
3007    }
3008}
3009
3010/// Future returned by [`WorkflowContext::sleep`].
3011pub struct TimerCall {
3012    ctx: WorkflowContext,
3013    delay_seconds: Option<u64>,
3014    scheduled: bool,
3015    matched_pending: bool,
3016}
3017
3018impl Future for TimerCall {
3019    type Output = Result<()>;
3020
3021    fn poll(mut self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
3022        if self.matched_pending {
3023            return Poll::Pending;
3024        }
3025
3026        let ctx = self.ctx.clone();
3027        let Some(requested_delay) = self.delay_seconds else {
3028            return Poll::Ready(Err(Error::TimerDurationOverflow));
3029        };
3030        let mut state = match ctx.state.lock() {
3031            Ok(state) => state,
3032            Err(_) => return Poll::Ready(Err(Error::WorkflowStatePoisoned)),
3033        };
3034
3035        if let Some(recorded) = state.recorded_commands.get(state.command_cursor).cloned() {
3036            match recorded {
3037                RecordedCommand::Timer {
3038                    sequence,
3039                    delay_seconds,
3040                    fired,
3041                    ..
3042                } => {
3043                    if delay_seconds != requested_delay {
3044                        return Poll::Ready(Err(Error::NonDeterministicReplay(
3045                            ReplayFailure::new(
3046                                "timer_delay_mismatch",
3047                                Some(sequence),
3048                                Some(format!("timer:{delay_seconds}s")),
3049                                Some(format!("timer:{requested_delay}s")),
3050                                "recorded timer delay differs from the current workflow command",
3051                            ),
3052                        )));
3053                    }
3054                    state.command_cursor += 1;
3055                    if fired {
3056                        return Poll::Ready(Ok(()));
3057                    }
3058                    state.matched_recorded_pending = true;
3059                    self.scheduled = true;
3060                    self.matched_pending = true;
3061                    return Poll::Pending;
3062                }
3063                other => return Poll::Ready(Err(command_mismatch(&other, "timer"))),
3064            }
3065        }
3066
3067        if !self.scheduled {
3068            state.commands.push(json!({
3069                "type": "start_timer",
3070                "delay_seconds": requested_delay,
3071            }));
3072            self.scheduled = true;
3073        }
3074
3075        Poll::Pending
3076    }
3077}
3078
3079/// Future returned by [`WorkflowContext::start_child_workflow`].
3080pub struct ChildWorkflowCall {
3081    ctx: WorkflowContext,
3082    workflow_type: String,
3083    options: ChildWorkflowOptions,
3084    args: Option<Result<Value>>,
3085    scheduled: bool,
3086    matched_pending: bool,
3087}
3088
3089impl Future for ChildWorkflowCall {
3090    type Output = Result<ChildWorkflowResult>;
3091
3092    fn poll(mut self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
3093        if self.matched_pending {
3094            return Poll::Pending;
3095        }
3096
3097        let ctx = self.ctx.clone();
3098        let mut state = match ctx.state.lock() {
3099            Ok(state) => state,
3100            Err(_) => return Poll::Ready(Err(Error::WorkflowStatePoisoned)),
3101        };
3102
3103        if let Some(recorded) = state.recorded_commands.get(state.command_cursor).cloned() {
3104            let sequence = recorded.sequence();
3105            match recorded {
3106                RecordedCommand::ChildWorkflow {
3107                    workflow_type,
3108                    outcome,
3109                    ..
3110                } => {
3111                    if let Some(recorded_type) = workflow_type {
3112                        if recorded_type != self.workflow_type {
3113                            return Poll::Ready(Err(Error::NonDeterministicReplay(
3114                                ReplayFailure::new(
3115                                    "recorded_command_detail_mismatch",
3116                                    Some(sequence),
3117                                    Some(format!("child workflow:{recorded_type}")),
3118                                    Some(format!("child workflow:{}", self.workflow_type)),
3119                                    "recorded child workflow type differs from the current workflow command",
3120                                ),
3121                            )));
3122                        }
3123                    }
3124                    state.command_cursor += 1;
3125                    if let Some(outcome) = outcome {
3126                        return Poll::Ready(outcome.map_err(Error::ChildWorkflowFailed));
3127                    }
3128                    state.matched_recorded_pending = true;
3129                    self.scheduled = true;
3130                    self.matched_pending = true;
3131                    return Poll::Pending;
3132                }
3133                other => {
3134                    return Poll::Ready(Err(command_mismatch(
3135                        &other,
3136                        format!("child workflow:{}", self.workflow_type),
3137                    )));
3138                }
3139            }
3140        }
3141
3142        if !self.scheduled {
3143            if self.options.task_queue.trim().is_empty() {
3144                return Poll::Ready(Err(Error::InvalidChildWorkflowOptions(
3145                    "task_queue must not be empty".to_string(),
3146                )));
3147            }
3148            for (name, value) in [
3149                (
3150                    "execution_timeout_seconds",
3151                    self.options.execution_timeout_seconds,
3152                ),
3153                ("run_timeout_seconds", self.options.run_timeout_seconds),
3154            ] {
3155                if value == Some(0) {
3156                    return Poll::Ready(Err(Error::InvalidChildWorkflowOptions(format!(
3157                        "{name} must be at least 1"
3158                    ))));
3159                }
3160            }
3161
3162            let args = match self.args.take().unwrap_or(Ok(Value::Null)) {
3163                Ok(args) => args,
3164                Err(error) => return Poll::Ready(Err(error)),
3165            };
3166            let arguments =
3167                match encode_value_envelope(&normalize_arguments(args), &state.payload_codec) {
3168                    Ok(arguments) => arguments,
3169                    Err(error) => return Poll::Ready(Err(error)),
3170                };
3171            let mut command = json!({
3172                "type": "start_child_workflow",
3173                "workflow_type": self.workflow_type,
3174                "queue": self.options.task_queue,
3175                "parent_close_policy": self.options.parent_close_policy.as_str(),
3176                "arguments": arguments,
3177            });
3178            let object = command
3179                .as_object_mut()
3180                .expect("child workflow command is always an object");
3181            if let Some(policy) = &self.options.retry_policy {
3182                let mut retry_policy = serde_json::Map::new();
3183                if let Some(max_attempts) = policy.max_attempts {
3184                    if max_attempts == 0 {
3185                        return Poll::Ready(Err(Error::InvalidChildWorkflowOptions(
3186                            "retry_policy.max_attempts must be at least 1".to_string(),
3187                        )));
3188                    }
3189                    retry_policy.insert("max_attempts".to_string(), json!(max_attempts));
3190                }
3191                if !policy.backoff_seconds.is_empty() {
3192                    retry_policy
3193                        .insert("backoff_seconds".to_string(), json!(policy.backoff_seconds));
3194                }
3195                if !policy.non_retryable_error_types.is_empty() {
3196                    retry_policy.insert(
3197                        "non_retryable_error_types".to_string(),
3198                        json!(policy.non_retryable_error_types),
3199                    );
3200                }
3201                if retry_policy.is_empty() {
3202                    return Poll::Ready(Err(Error::InvalidChildWorkflowOptions(
3203                        "retry_policy must configure at least one field".to_string(),
3204                    )));
3205                }
3206                object.insert("retry_policy".to_string(), Value::Object(retry_policy));
3207            }
3208            if let Some(seconds) = self.options.execution_timeout_seconds {
3209                object.insert("execution_timeout_seconds".to_string(), json!(seconds));
3210            }
3211            if let Some(seconds) = self.options.run_timeout_seconds {
3212                object.insert("run_timeout_seconds".to_string(), json!(seconds));
3213            }
3214            state.commands.push(command);
3215            self.scheduled = true;
3216        }
3217
3218        Poll::Pending
3219    }
3220}
3221
3222fn command_mismatch(recorded: &RecordedCommand, actual: impl Into<String>) -> Error {
3223    Error::NonDeterministicReplay(ReplayFailure::new(
3224        "recorded_command_mismatch",
3225        Some(recorded.sequence()),
3226        Some(recorded.shape().to_string()),
3227        Some(actual.into()),
3228        "current workflow command does not match the recorded durable command sequence",
3229    ))
3230}
3231
3232pub struct SignalCall {
3233    ctx: WorkflowContext,
3234    signal_name: String,
3235    opened_wait: bool,
3236    matched_pending: bool,
3237}
3238
3239impl Future for SignalCall {
3240    type Output = Result<Vec<Value>>;
3241
3242    fn poll(mut self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
3243        if self.matched_pending {
3244            return Poll::Pending;
3245        }
3246
3247        let ctx = self.ctx.clone();
3248        let mut state = match ctx.state.lock() {
3249            Ok(state) => state,
3250            Err(_) => return Poll::Ready(Err(Error::WorkflowStatePoisoned)),
3251        };
3252
3253        let signals = match signal_values(
3254            &state.history,
3255            &self.signal_name,
3256            &state.payload_codec,
3257            state.resume_signal.as_ref(),
3258        ) {
3259            Ok(signals) => signals,
3260            Err(error) => return Poll::Ready(Err(error)),
3261        };
3262        let cursor = *state.signal_cursors.get(&self.signal_name).unwrap_or(&0);
3263
3264        if let Some(recorded) = state.recorded_commands.get(state.command_cursor).cloned() {
3265            match recorded {
3266                RecordedCommand::SignalWait {
3267                    sequence,
3268                    signal_name,
3269                } => {
3270                    if let Some(recorded_name) = signal_name {
3271                        if recorded_name != self.signal_name {
3272                            return Poll::Ready(Err(Error::NonDeterministicReplay(
3273                                ReplayFailure::new(
3274                                    "recorded_command_detail_mismatch",
3275                                    Some(sequence),
3276                                    Some(format!("signal wait:{recorded_name}")),
3277                                    Some(format!("signal wait:{}", self.signal_name)),
3278                                    "recorded signal name differs from the current workflow command",
3279                                ),
3280                            )));
3281                        }
3282                    }
3283
3284                    state.command_cursor += 1;
3285                    if cursor < signals.len() {
3286                        state
3287                            .signal_cursors
3288                            .insert(self.signal_name.clone(), cursor + 1);
3289                        return Poll::Ready(Ok(signals[cursor].clone()));
3290                    }
3291
3292                    state.matched_recorded_pending = true;
3293                    self.opened_wait = true;
3294                    self.matched_pending = true;
3295                    return Poll::Pending;
3296                }
3297                other => {
3298                    return Poll::Ready(Err(command_mismatch(
3299                        &other,
3300                        format!("signal wait:{}", self.signal_name),
3301                    )));
3302                }
3303            }
3304        }
3305
3306        if cursor < signals.len() {
3307            state
3308                .signal_cursors
3309                .insert(self.signal_name.clone(), cursor + 1);
3310            return Poll::Ready(Ok(signals[cursor].clone()));
3311        }
3312
3313        if !self.opened_wait {
3314            state.commands.push(json!({
3315                "type": "open_condition_wait",
3316                "condition_key": format!("signal:{}", self.signal_name)
3317            }));
3318            self.opened_wait = true;
3319        }
3320
3321        Poll::Pending
3322    }
3323}
3324
3325#[derive(Clone, Debug)]
3326pub struct ActivityContext {
3327    client: Client,
3328    pub task_id: String,
3329    pub activity_attempt_id: String,
3330    pub lease_owner: String,
3331    pub activity_type: String,
3332    pub attempt_number: u64,
3333    pub task_queue: String,
3334    pub worker_id: String,
3335}
3336
3337impl ActivityContext {
3338    pub async fn heartbeat<T: Serialize>(&self, details: T) -> Result<ActivityHeartbeatResponse> {
3339        self.client
3340            .heartbeat_activity_task(
3341                &self.task_id,
3342                &self.activity_attempt_id,
3343                &self.lease_owner,
3344                serde_json::to_value(details)?,
3345            )
3346            .await
3347    }
3348}
3349
3350fn decode_task_arguments(value: Option<&Value>, codec: &str) -> Result<Value> {
3351    match value {
3352        Some(value) => Ok(normalize_arguments(decode_wire_value(value, codec)?)),
3353        None => Ok(Value::Array(Vec::new())),
3354    }
3355}
3356
3357fn decode_resume_signal(task: &WorkflowTask) -> Result<Option<ResumeSignal>> {
3358    let Some(signal_name) = task
3359        .signal_name
3360        .as_deref()
3361        .filter(|value| !value.is_empty())
3362    else {
3363        return Ok(None);
3364    };
3365    let Some(arguments) = task.signal_arguments.as_ref() else {
3366        return Ok(None);
3367    };
3368
3369    let decoded = normalize_arguments(decode_wire_value(arguments, &task.payload_codec)?);
3370    let Value::Array(arguments) = decoded else {
3371        unreachable!("normalize_arguments always returns an array");
3372    };
3373
3374    Ok(Some(ResumeSignal {
3375        signal_id: task.workflow_signal_id.clone(),
3376        signal_name: signal_name.to_string(),
3377        arguments,
3378    }))
3379}
3380
3381fn normalize_arguments(value: Value) -> Value {
3382    match value {
3383        Value::Null => Value::Array(Vec::new()),
3384        Value::Array(_) => value,
3385        other => Value::Array(vec![other]),
3386    }
3387}
3388
3389fn recorded_commands(
3390    events: &[HistoryEvent],
3391    fallback_codec: &str,
3392    parent: WorkflowIdentity,
3393) -> Result<Vec<RecordedCommand>> {
3394    let mut events_by_sequence: BTreeMap<u64, Vec<&HistoryEvent>> = BTreeMap::new();
3395
3396    for event in events {
3397        let is_activity = matches!(
3398            event.event_type.as_str(),
3399            "ActivityScheduled"
3400                | "ActivityStarted"
3401                | "ActivityHeartbeatRecorded"
3402                | "ActivityRetryScheduled"
3403                | "ActivityCompleted"
3404                | "ActivityFailed"
3405                | "ActivityCancelled"
3406                | "ActivityTimedOut"
3407        );
3408        let is_workflow_timer = matches!(
3409            event.event_type.as_str(),
3410            "TimerScheduled" | "TimerCancelled" | "TimerFired"
3411        ) && !is_internal_timer_event(event);
3412        let is_child_workflow = matches!(
3413            event.event_type.as_str(),
3414            "ChildWorkflowScheduled"
3415                | "ChildRunCompleted"
3416                | "ChildRunFailed"
3417                | "ChildRunCancelled"
3418                | "ChildRunTerminated"
3419        );
3420        let is_signal_wait = is_recorded_signal_wait_event(event);
3421        if !is_activity && !is_workflow_timer && !is_child_workflow && !is_signal_wait {
3422            continue;
3423        }
3424
3425        let sequence = durable_event_sequence(event).ok_or_else(|| {
3426            Error::NonDeterministicReplay(ReplayFailure::new(
3427                "durable_command_sequence_missing",
3428                None,
3429                Some("positive workflow sequence".to_string()),
3430                Some(event.event_type.clone()),
3431                "durable command history event has no workflow sequence",
3432            ))
3433        })?;
3434        if sequence == 0 {
3435            return Err(Error::NonDeterministicReplay(ReplayFailure::new(
3436                "durable_command_sequence_invalid",
3437                Some(sequence),
3438                Some("positive workflow sequence".to_string()),
3439                Some(sequence.to_string()),
3440                "durable command history uses an invalid workflow sequence",
3441            )));
3442        }
3443        events_by_sequence.entry(sequence).or_default().push(event);
3444    }
3445
3446    events_by_sequence
3447        .into_iter()
3448        .map(|(sequence, sequence_events)| {
3449            let activity_events: Vec<_> = sequence_events
3450                .iter()
3451                .copied()
3452                .filter(|event| event.event_type.starts_with("Activity"))
3453                .collect();
3454            let timer_events: Vec<_> = sequence_events
3455                .iter()
3456                .copied()
3457                .filter(|event| event.event_type.starts_with("Timer"))
3458                .collect();
3459            let child_events: Vec<_> = sequence_events
3460                .iter()
3461                .copied()
3462                .filter(|event| {
3463                    event.event_type == "ChildWorkflowScheduled"
3464                        || event.event_type.starts_with("ChildRun")
3465                })
3466                .collect();
3467            let signal_wait_events: Vec<_> = sequence_events
3468                .iter()
3469                .copied()
3470                .filter(|event| is_recorded_signal_wait_event(event))
3471                .collect();
3472
3473            let command_kind_count = usize::from(!activity_events.is_empty())
3474                + usize::from(!timer_events.is_empty())
3475                + usize::from(!child_events.is_empty())
3476                + usize::from(!signal_wait_events.is_empty());
3477            if command_kind_count > 1 {
3478                let actual = [
3479                    (!activity_events.is_empty()).then_some("activity"),
3480                    (!timer_events.is_empty()).then_some("timer"),
3481                    (!child_events.is_empty()).then_some("child workflow"),
3482                    (!signal_wait_events.is_empty()).then_some("signal wait"),
3483                ]
3484                .into_iter()
3485                .flatten()
3486                .collect::<Vec<_>>()
3487                .join(" and ");
3488                return Err(invalid_recorded_history(
3489                    "durable_command_sequence_collision",
3490                    sequence,
3491                    "one durable command kind",
3492                    &actual,
3493                    "one workflow sequence records more than one durable command kind",
3494                ));
3495            }
3496
3497            if !activity_events.is_empty() {
3498                let scheduled_count = activity_events
3499                    .iter()
3500                    .filter(|event| event.event_type == "ActivityScheduled")
3501                    .count();
3502                if scheduled_count > 1 {
3503                    return Err(invalid_recorded_history(
3504                        "duplicate_activity_schedule",
3505                        sequence,
3506                        "at most one ActivityScheduled event",
3507                        "multiple ActivityScheduled events",
3508                        "activity history schedules more than one command at one workflow sequence",
3509                    ));
3510                }
3511                let activity_type = activity_events.iter().find_map(|event| {
3512                    event
3513                        .payload
3514                        .get("activity_type")
3515                        .or_else(|| event.payload.get("activity_name"))
3516                        .and_then(Value::as_str)
3517                        .map(str::to_string)
3518                });
3519                if activity_events.iter().filter_map(|event| {
3520                    event
3521                        .payload
3522                        .get("activity_type")
3523                        .or_else(|| event.payload.get("activity_name"))
3524                        .and_then(Value::as_str)
3525                }).any(|candidate| Some(candidate) != activity_type.as_deref()) {
3526                    return Err(invalid_recorded_history(
3527                        "activity_identity_mismatch",
3528                        sequence,
3529                        activity_type.as_deref().unwrap_or("one activity identity"),
3530                        "conflicting activity identities",
3531                        "activity lifecycle events at one workflow sequence disagree on identity",
3532                    ));
3533                }
3534                let completed: Vec<_> = activity_events
3535                    .iter()
3536                    .copied()
3537                    .filter(|event| event.event_type == "ActivityCompleted")
3538                    .collect();
3539                if completed.len() > 1 {
3540                    return Err(invalid_recorded_history(
3541                        "duplicate_activity_completion",
3542                        sequence,
3543                        "one ActivityCompleted event",
3544                        "multiple ActivityCompleted events",
3545                        "activity history contains more than one completion for a workflow sequence",
3546                    ));
3547                }
3548                let result = completed
3549                    .first()
3550                    .map(|event| {
3551                        let codec = event
3552                            .payload
3553                            .get("payload_codec")
3554                            .and_then(Value::as_str)
3555                            .unwrap_or(fallback_codec);
3556                        decode_wire_value(
3557                            event.payload.get("result").unwrap_or(&Value::Null),
3558                            codec,
3559                        )
3560                    })
3561                    .transpose()?;
3562                return Ok(RecordedCommand::Activity {
3563                    sequence,
3564                    activity_type,
3565                    result,
3566                });
3567            }
3568
3569            if !child_events.is_empty() {
3570                let scheduled: Vec<_> = child_events
3571                    .iter()
3572                    .copied()
3573                    .filter(|event| event.event_type == "ChildWorkflowScheduled")
3574                    .collect();
3575                if scheduled.len() != 1 {
3576                    return Err(invalid_recorded_history(
3577                        "child_workflow_schedule_missing_or_duplicate",
3578                        sequence,
3579                        "one ChildWorkflowScheduled event",
3580                        &format!("{} ChildWorkflowScheduled events", scheduled.len()),
3581                        "child workflow replay requires exactly one recorded schedule event",
3582                    ));
3583                }
3584                let workflow_type = child_events.iter().find_map(|event| {
3585                    event
3586                        .payload
3587                        .get("child_workflow_type")
3588                        .or_else(|| event.payload.get("workflow_type"))
3589                        .and_then(Value::as_str)
3590                        .filter(|value| !value.is_empty())
3591                        .map(str::to_string)
3592                });
3593                if child_events
3594                    .iter()
3595                    .filter_map(|event| {
3596                        event
3597                            .payload
3598                            .get("child_workflow_type")
3599                            .or_else(|| event.payload.get("workflow_type"))
3600                            .and_then(Value::as_str)
3601                    })
3602                    .any(|candidate| Some(candidate) != workflow_type.as_deref())
3603                {
3604                    return Err(invalid_recorded_history(
3605                        "child_workflow_identity_mismatch",
3606                        sequence,
3607                        workflow_type
3608                            .as_deref()
3609                            .unwrap_or("one child workflow type"),
3610                        "conflicting child workflow types",
3611                        "child workflow lifecycle events at one sequence disagree on type",
3612                    ));
3613                }
3614                let mut outcomes = child_workflow_outcomes(
3615                    &child_events.iter().map(|event| (*event).clone()).collect::<Vec<_>>(),
3616                    fallback_codec,
3617                    parent.clone(),
3618                )?;
3619                if outcomes.len() > 1 {
3620                    return Err(invalid_recorded_history(
3621                        "duplicate_child_workflow_terminal_event",
3622                        sequence,
3623                        "at most one terminal child event",
3624                        "multiple terminal child events",
3625                        "child workflow history settles one command more than once",
3626                    ));
3627                }
3628                return Ok(RecordedCommand::ChildWorkflow {
3629                    sequence,
3630                    workflow_type,
3631                    outcome: outcomes.pop(),
3632                });
3633            }
3634
3635            if !signal_wait_events.is_empty() {
3636                let opened_count = signal_wait_events
3637                    .iter()
3638                    .filter(|event| {
3639                        matches!(
3640                            event.event_type.as_str(),
3641                            "SignalWaitOpened" | "ConditionWaitOpened"
3642                        )
3643                    })
3644                    .count();
3645                if opened_count > 1 {
3646                    return Err(invalid_recorded_history(
3647                        "duplicate_signal_wait_open",
3648                        sequence,
3649                        "at most one signal wait open event",
3650                        "multiple signal wait open events",
3651                        "signal history opens more than one durable wait at one workflow sequence",
3652                    ));
3653                }
3654
3655                let signal_name = signal_wait_events
3656                    .iter()
3657                    .find_map(|event| recorded_signal_wait_name(event));
3658                if signal_wait_events
3659                    .iter()
3660                    .filter_map(|event| recorded_signal_wait_name(event))
3661                    .any(|candidate| Some(candidate.as_str()) != signal_name.as_deref())
3662                {
3663                    return Err(invalid_recorded_history(
3664                        "signal_wait_identity_mismatch",
3665                        sequence,
3666                        signal_name.as_deref().unwrap_or("one signal name"),
3667                        "conflicting signal names",
3668                        "signal wait lifecycle events at one workflow sequence disagree on identity",
3669                    ));
3670                }
3671                return Ok(RecordedCommand::SignalWait {
3672                    sequence,
3673                    signal_name,
3674                });
3675            }
3676
3677            let scheduled: Vec<_> = timer_events
3678                .iter()
3679                .copied()
3680                .filter(|event| event.event_type == "TimerScheduled")
3681                .collect();
3682            let fired: Vec<_> = timer_events
3683                .iter()
3684                .copied()
3685                .filter(|event| event.event_type == "TimerFired")
3686                .collect();
3687            if scheduled.len() != 1 {
3688                return Err(invalid_recorded_history(
3689                    "timer_schedule_missing_or_duplicate",
3690                    sequence,
3691                    "one TimerScheduled event",
3692                    &format!("{} TimerScheduled events", scheduled.len()),
3693                    "timer replay requires exactly one recorded schedule event",
3694                ));
3695            }
3696            if fired.len() > 1 {
3697                return Err(invalid_recorded_history(
3698                    "duplicate_timer_fire",
3699                    sequence,
3700                    "at most one TimerFired event",
3701                    "multiple TimerFired events",
3702                    "timer history contains more than one fire event for a workflow sequence",
3703                ));
3704            }
3705
3706            let scheduled = scheduled[0];
3707            let timer_id = required_history_string(scheduled, "timer_id", sequence)?;
3708            let delay_seconds = required_history_u64(scheduled, "delay_seconds", sequence)?;
3709            if let Some(fired) = fired.first() {
3710                let fired_timer_id = required_history_string(fired, "timer_id", sequence)?;
3711                if fired_timer_id != timer_id {
3712                    return Err(invalid_recorded_history(
3713                        "timer_identity_mismatch",
3714                        sequence,
3715                        &timer_id,
3716                        &fired_timer_id,
3717                        "TimerFired does not correspond to the recorded TimerScheduled event",
3718                    ));
3719                }
3720                let fired_delay = required_history_u64(fired, "delay_seconds", sequence)?;
3721                if fired_delay != delay_seconds {
3722                    return Err(invalid_recorded_history(
3723                        "timer_history_delay_mismatch",
3724                        sequence,
3725                        &delay_seconds.to_string(),
3726                        &fired_delay.to_string(),
3727                        "TimerScheduled and TimerFired record different delays",
3728                    ));
3729                }
3730            }
3731
3732            Ok(RecordedCommand::Timer {
3733                sequence,
3734                delay_seconds,
3735                fired: !fired.is_empty(),
3736            })
3737        })
3738        .collect()
3739}
3740
3741fn durable_event_sequence(event: &HistoryEvent) -> Option<u64> {
3742    event
3743        .payload
3744        .get("sequence")
3745        .or_else(|| event.payload.get("workflow_sequence"))
3746        .or_else(|| event.raw.get("sequence"))
3747        .or_else(|| event.raw.get("workflow_sequence"))
3748        .and_then(value_as_u64)
3749}
3750
3751fn is_internal_timer_event(event: &HistoryEvent) -> bool {
3752    matches!(
3753        event
3754            .payload
3755            .get("timer_kind")
3756            .or_else(|| event.raw.get("timer_kind"))
3757            .and_then(Value::as_str),
3758        Some("condition_timeout" | "signal_timeout")
3759    )
3760}
3761
3762fn recorded_signal_wait_name(event: &HistoryEvent) -> Option<String> {
3763    match event.event_type.as_str() {
3764        "SignalWaitOpened" | "SignalApplied" => event
3765            .payload
3766            .get("signal_name")
3767            .or_else(|| event.raw.get("signal_name"))
3768            .and_then(Value::as_str)
3769            .filter(|value| !value.is_empty())
3770            .map(str::to_string),
3771        "ConditionWaitOpened" | "ConditionWaitSatisfied" | "ConditionWaitTimedOut" => event
3772            .payload
3773            .get("condition_key")
3774            .or_else(|| event.raw.get("condition_key"))
3775            .and_then(Value::as_str)
3776            .and_then(|key| key.strip_prefix("signal:"))
3777            .filter(|value| !value.is_empty())
3778            .map(str::to_string),
3779        _ => None,
3780    }
3781}
3782
3783fn is_recorded_signal_wait_event(event: &HistoryEvent) -> bool {
3784    match event.event_type.as_str() {
3785        "SignalWaitOpened" | "SignalApplied" => true,
3786        "ConditionWaitOpened" | "ConditionWaitSatisfied" | "ConditionWaitTimedOut" => event
3787            .payload
3788            .get("condition_key")
3789            .or_else(|| event.raw.get("condition_key"))
3790            .and_then(Value::as_str)
3791            .is_some_and(|key| key.starts_with("signal:")),
3792        _ => false,
3793    }
3794}
3795
3796fn required_history_string(event: &HistoryEvent, field: &str, sequence: u64) -> Result<String> {
3797    event
3798        .payload
3799        .get(field)
3800        .and_then(Value::as_str)
3801        .filter(|value| !value.is_empty())
3802        .map(str::to_string)
3803        .ok_or_else(|| {
3804            invalid_recorded_history(
3805                "timer_history_field_missing",
3806                sequence,
3807                field,
3808                &event.event_type,
3809                "timer history is missing a required identity field",
3810            )
3811        })
3812}
3813
3814fn required_history_u64(event: &HistoryEvent, field: &str, sequence: u64) -> Result<u64> {
3815    event
3816        .payload
3817        .get(field)
3818        .and_then(value_as_u64)
3819        .ok_or_else(|| {
3820            invalid_recorded_history(
3821                "timer_history_field_missing",
3822                sequence,
3823                field,
3824                &event.event_type,
3825                "timer history is missing a required numeric field",
3826            )
3827        })
3828}
3829
3830fn invalid_recorded_history(
3831    reason: &str,
3832    sequence: u64,
3833    expected: &str,
3834    actual: &str,
3835    message: &str,
3836) -> Error {
3837    Error::NonDeterministicReplay(ReplayFailure::new(
3838        reason,
3839        Some(sequence),
3840        Some(expected.to_string()),
3841        Some(actual.to_string()),
3842        message,
3843    ))
3844}
3845
3846type ChildWorkflowOutcome = std::result::Result<ChildWorkflowResult, ChildWorkflowFailure>;
3847
3848fn child_workflow_outcomes(
3849    events: &[HistoryEvent],
3850    fallback_codec: &str,
3851    parent: WorkflowIdentity,
3852) -> Result<Vec<ChildWorkflowOutcome>> {
3853    let mut outcomes = Vec::new();
3854
3855    for event in events {
3856        let kind = match event.event_type.as_str() {
3857            "ChildRunCompleted" => None,
3858            "ChildRunFailed" => Some((
3859                ChildWorkflowFailureKind::Failed,
3860                "child_workflow",
3861                "child workflow failed",
3862            )),
3863            "ChildRunCancelled" => Some((
3864                ChildWorkflowFailureKind::Cancelled,
3865                "cancelled",
3866                "child workflow was cancelled",
3867            )),
3868            "ChildRunTerminated" => Some((
3869                ChildWorkflowFailureKind::Terminated,
3870                "terminated",
3871                "child workflow was terminated",
3872            )),
3873            _ => continue,
3874        };
3875        let payload = &event.payload;
3876        let child_workflow_id = payload_string(payload, "child_workflow_instance_id");
3877        let child_workflow_run_id = payload_string(payload, "child_workflow_run_id");
3878        let child_workflow_type = payload_string(payload, "child_workflow_type");
3879
3880        if let Some((kind, reason, fallback_message)) = kind {
3881            let exception = payload
3882                .get("exception")
3883                .filter(|value| !value.is_null())
3884                .cloned();
3885            let message = payload_string(payload, "message")
3886                .or_else(|| {
3887                    exception
3888                        .as_ref()
3889                        .and_then(|value| payload_string(value, "message"))
3890                })
3891                .unwrap_or_else(|| fallback_message.to_string());
3892            let exception_type = payload_string(payload, "exception_type").or_else(|| {
3893                exception
3894                    .as_ref()
3895                    .and_then(|value| payload_string(value, "type"))
3896            });
3897            let exception_class = payload_string(payload, "exception_class").or_else(|| {
3898                exception
3899                    .as_ref()
3900                    .and_then(|value| payload_string(value, "class"))
3901            });
3902            outcomes.push(Err(ChildWorkflowFailure {
3903                kind,
3904                reason: reason.to_string(),
3905                message,
3906                parent_workflow_id: parent.workflow_id.clone(),
3907                parent_workflow_run_id: parent.run_id.clone(),
3908                child_workflow_id,
3909                child_workflow_run_id,
3910                child_workflow_type,
3911                failure_id: payload_string(payload, "failure_id"),
3912                failure_category: payload_string(payload, "failure_category"),
3913                exception_type,
3914                exception_class,
3915                non_retryable: payload
3916                    .get("non_retryable")
3917                    .and_then(Value::as_bool)
3918                    .unwrap_or(false),
3919                code: payload
3920                    .get("code")
3921                    .filter(|value| !value.is_null())
3922                    .cloned(),
3923                exception,
3924            }));
3925            continue;
3926        }
3927
3928        let codec = payload
3929            .get("payload_codec")
3930            .and_then(Value::as_str)
3931            .unwrap_or(fallback_codec);
3932        let result = payload
3933            .get("result")
3934            .or_else(|| payload.get("output"))
3935            .unwrap_or(&Value::Null);
3936        outcomes.push(Ok(ChildWorkflowResult {
3937            parent: parent.clone(),
3938            child: WorkflowIdentity {
3939                workflow_id: child_workflow_id,
3940                run_id: child_workflow_run_id,
3941            },
3942            child_workflow_type,
3943            result: decode_wire_value(result, codec)?,
3944        }));
3945    }
3946
3947    Ok(outcomes)
3948}
3949
3950fn payload_string(payload: &Value, key: &str) -> Option<String> {
3951    payload
3952        .get(key)
3953        .and_then(Value::as_str)
3954        .filter(|value| !value.is_empty())
3955        .map(str::to_string)
3956}
3957
3958fn workflow_failure_command(error: &Error) -> Value {
3959    let (exception_type, exception_class, properties) = match error {
3960        Error::ChildWorkflowFailed(failure) => (
3961            match failure.kind {
3962                ChildWorkflowFailureKind::Failed => "ChildWorkflowFailed",
3963                ChildWorkflowFailureKind::Cancelled => "ChildWorkflowCancelled",
3964                ChildWorkflowFailureKind::Terminated => "ChildWorkflowTerminated",
3965            },
3966            "durable_workflow::ChildWorkflowFailure",
3967            json!({
3968                "reason": failure.reason,
3969                "parent_workflow_id": failure.parent_workflow_id,
3970                "parent_workflow_run_id": failure.parent_workflow_run_id,
3971                "child_workflow_id": failure.child_workflow_id,
3972                "child_workflow_run_id": failure.child_workflow_run_id,
3973                "child_workflow_type": failure.child_workflow_type,
3974                "failure_id": failure.failure_id,
3975                "failure_category": failure.failure_category,
3976                "child_exception_type": failure.exception_type,
3977                "child_exception_class": failure.exception_class,
3978                "child_non_retryable": failure.non_retryable,
3979                "child_code": failure.code,
3980                "child_exception": failure.exception,
3981            }),
3982        ),
3983        Error::NonDeterministicReplay(_) => (
3984            "NonDeterministicReplay",
3985            "durable_workflow::Error",
3986            Value::Null,
3987        ),
3988        _ => ("RustWorkflowError", "durable_workflow::Error", Value::Null),
3989    };
3990    let non_retryable = match error {
3991        Error::ChildWorkflowFailed(failure) => failure.non_retryable,
3992        Error::NonDeterministicReplay(_) => true,
3993        _ => false,
3994    };
3995
3996    json!({
3997        "type": "fail_workflow",
3998        "message": error.to_string(),
3999        "exception_type": exception_type,
4000        "exception_class": exception_class,
4001        "non_retryable": non_retryable,
4002        "exception": {
4003            "type": exception_type,
4004            "class": exception_class,
4005            "message": error.to_string(),
4006            "properties": properties,
4007        }
4008    })
4009}
4010
4011fn signal_values(
4012    events: &[HistoryEvent],
4013    signal_name: &str,
4014    fallback_codec: &str,
4015    resume_signal: Option<&ResumeSignal>,
4016) -> Result<Vec<Vec<Value>>> {
4017    let mut signals = Vec::new();
4018
4019    for event in events {
4020        if event.event_type != "SignalApplied" && event.event_type != "SignalReceived" {
4021            continue;
4022        }
4023
4024        if event.payload.get("signal_name").and_then(Value::as_str) != Some(signal_name) {
4025            continue;
4026        }
4027
4028        let codec = event
4029            .payload
4030            .get("payload_codec")
4031            .and_then(Value::as_str)
4032            .unwrap_or(fallback_codec);
4033        let raw = event
4034            .payload
4035            .get("value")
4036            .or_else(|| event.payload.get("input"))
4037            .or_else(|| event.payload.get("arguments"));
4038        let decoded = match raw.filter(|value| !value.is_null()) {
4039            Some(value) => decode_wire_value(value, codec)?,
4040            None => resume_signal
4041                .filter(|signal| resume_signal_matches_event(signal, event, signal_name))
4042                .map(|signal| Value::Array(signal.arguments.clone()))
4043                .unwrap_or_else(|| Value::Array(Vec::new())),
4044        };
4045        let args = match normalize_arguments(decoded) {
4046            Value::Array(values) => values,
4047            _ => unreachable!("normalize_arguments always returns an array"),
4048        };
4049        signals.push(args);
4050    }
4051
4052    Ok(signals)
4053}
4054
4055fn hydrate_query_history_from_export(task: &mut QueryTask) -> Result<()> {
4056    let Some(export_events) = task
4057        .history_export
4058        .as_ref()
4059        .and_then(|export| export.get("history_events"))
4060        .and_then(Value::as_array)
4061    else {
4062        return Ok(());
4063    };
4064
4065    if export_events.len() > task.history_events.len() {
4066        task.history_events = serde_json::from_value(Value::Array(export_events.clone()))?;
4067    }
4068
4069    Ok(())
4070}
4071
4072fn enrich_query_history_from_export(task: &mut QueryTask) -> Result<()> {
4073    let Some(export) = task.history_export.as_ref() else {
4074        return Ok(());
4075    };
4076    let signals = export
4077        .get("signals")
4078        .and_then(Value::as_array)
4079        .cloned()
4080        .unwrap_or_default();
4081    let activities = export
4082        .get("activities")
4083        .and_then(Value::as_array)
4084        .cloned()
4085        .unwrap_or_default();
4086    let export_codec = export
4087        .get("payloads")
4088        .and_then(|payloads| payloads.get("codec"))
4089        .and_then(Value::as_str)
4090        .unwrap_or(&task.payload_codec)
4091        .to_string();
4092    let mut signal_name_offsets: HashMap<String, usize> = HashMap::new();
4093
4094    for event in &mut task.history_events {
4095        if event.event_type == "ActivityCompleted" {
4096            let sequence = event
4097                .payload
4098                .get("sequence")
4099                .or_else(|| event.payload.get("workflow_sequence"))
4100                .and_then(value_as_u64);
4101            let Some(activity) = sequence.and_then(|sequence| {
4102                activities.iter().find(|activity| {
4103                    activity.get("sequence").and_then(value_as_u64) == Some(sequence)
4104                })
4105            }) else {
4106                continue;
4107            };
4108            let Some(payload) = event.payload.as_object_mut() else {
4109                continue;
4110            };
4111            if missing_payload(payload.get("result")) {
4112                if let Some(result) = activity
4113                    .get("result")
4114                    .filter(|value| !missing_payload(Some(value)))
4115                {
4116                    payload.insert("result".to_string(), result.clone());
4117                }
4118            }
4119            for field in ["payload_codec", "activity_type"] {
4120                if payload
4121                    .get(field)
4122                    .and_then(Value::as_str)
4123                    .unwrap_or_default()
4124                    .is_empty()
4125                {
4126                    if let Some(value) = activity.get(field) {
4127                        payload.insert(field.to_string(), value.clone());
4128                    }
4129                }
4130            }
4131            continue;
4132        }
4133
4134        if event.event_type != "SignalReceived" && event.event_type != "SignalApplied" {
4135            continue;
4136        }
4137        let signal_id = event.payload.get("signal_id").and_then(Value::as_str);
4138        let command_id = event
4139            .payload
4140            .get("workflow_command_id")
4141            .or_else(|| event.raw.get("workflow_command_id"))
4142            .and_then(Value::as_str);
4143        let signal_name = event
4144            .payload
4145            .get("signal_name")
4146            .and_then(Value::as_str)
4147            .unwrap_or_default()
4148            .to_string();
4149        let matched = signals
4150            .iter()
4151            .find(|signal| {
4152                signal_id.is_some() && signal.get("id").and_then(Value::as_str) == signal_id
4153            })
4154            .or_else(|| {
4155                signals.iter().find(|signal| {
4156                    command_id.is_some()
4157                        && signal.get("command_id").and_then(Value::as_str) == command_id
4158                })
4159            })
4160            .or_else(|| {
4161                let offset = signal_name_offsets.entry(signal_name.clone()).or_default();
4162                let signal = signals
4163                    .iter()
4164                    .filter(|signal| {
4165                        signal.get("name").and_then(Value::as_str) == Some(signal_name.as_str())
4166                    })
4167                    .nth(*offset);
4168                if signal.is_some() {
4169                    *offset += 1;
4170                }
4171                signal
4172            });
4173        let Some(signal) = matched else {
4174            continue;
4175        };
4176        let signal_codec = signal
4177            .get("payload_codec")
4178            .and_then(Value::as_str)
4179            .unwrap_or(&export_codec);
4180        let Some(payload) = event.payload.as_object_mut() else {
4181            continue;
4182        };
4183        if missing_payload(payload.get("arguments")) {
4184            if let Some(arguments) = signal
4185                .get("arguments")
4186                .filter(|value| !missing_payload(Some(value)))
4187            {
4188                let envelope = match arguments {
4189                    Value::String(blob) => json!({"codec": signal_codec, "blob": blob}),
4190                    other => other.clone(),
4191                };
4192                payload.insert("arguments".to_string(), envelope);
4193            }
4194        }
4195        if payload
4196            .get("payload_codec")
4197            .and_then(Value::as_str)
4198            .unwrap_or_default()
4199            .is_empty()
4200        {
4201            payload.insert("payload_codec".to_string(), json!(signal_codec));
4202        }
4203    }
4204
4205    Ok(())
4206}
4207
4208fn missing_payload(value: Option<&Value>) -> bool {
4209    match value {
4210        None | Some(Value::Null) => true,
4211        Some(Value::String(value)) => value.is_empty(),
4212        Some(_) => false,
4213    }
4214}
4215
4216fn query_signal_events(task: &QueryTask) -> Result<Vec<QuerySignal>> {
4217    let export_signals = task
4218        .history_export
4219        .as_ref()
4220        .and_then(|export| export.get("signals"))
4221        .and_then(Value::as_array)
4222        .cloned()
4223        .unwrap_or_default();
4224    let export_codec = task
4225        .history_export
4226        .as_ref()
4227        .and_then(|export| export.get("payloads"))
4228        .and_then(|payloads| payloads.get("codec"))
4229        .and_then(Value::as_str)
4230        .unwrap_or(&task.payload_codec);
4231    let mut name_offsets: HashMap<String, usize> = HashMap::new();
4232    let mut signals = Vec::new();
4233
4234    for event in &task.history_events {
4235        if event.event_type != "SignalApplied" && event.event_type != "SignalReceived" {
4236            continue;
4237        }
4238
4239        let name = event
4240            .payload
4241            .get("signal_name")
4242            .and_then(Value::as_str)
4243            .unwrap_or_default();
4244        if name.is_empty() {
4245            continue;
4246        }
4247        let signal_id = event.payload.get("signal_id").and_then(Value::as_str);
4248        let command_id = event
4249            .payload
4250            .get("workflow_command_id")
4251            .or_else(|| event.raw.get("workflow_command_id"))
4252            .and_then(Value::as_str);
4253        let matched_export = export_signals
4254            .iter()
4255            .find(|candidate| {
4256                signal_id.is_some() && candidate.get("id").and_then(Value::as_str) == signal_id
4257            })
4258            .or_else(|| {
4259                export_signals.iter().find(|candidate| {
4260                    command_id.is_some()
4261                        && candidate.get("command_id").and_then(Value::as_str) == command_id
4262                })
4263            })
4264            .or_else(|| {
4265                let offset = name_offsets.entry(name.to_string()).or_default();
4266                let candidate = export_signals
4267                    .iter()
4268                    .filter(|candidate| candidate.get("name").and_then(Value::as_str) == Some(name))
4269                    .nth(*offset);
4270                if candidate.is_some() {
4271                    *offset += 1;
4272                }
4273                candidate
4274            });
4275        let codec = event
4276            .payload
4277            .get("payload_codec")
4278            .and_then(Value::as_str)
4279            .or_else(|| {
4280                matched_export
4281                    .and_then(|signal| signal.get("payload_codec"))
4282                    .and_then(Value::as_str)
4283            })
4284            .unwrap_or(export_codec);
4285        let raw_arguments = event
4286            .payload
4287            .get("value")
4288            .or_else(|| event.payload.get("input"))
4289            .or_else(|| event.payload.get("arguments"))
4290            .filter(|value| !value.is_null())
4291            .or_else(|| matched_export.and_then(|signal| signal.get("arguments")));
4292        let arguments = decode_query_signal_arguments(raw_arguments, codec)?;
4293        let workflow_sequence = event
4294            .payload
4295            .get("workflow_sequence")
4296            .and_then(value_as_u64)
4297            .or_else(|| {
4298                matched_export
4299                    .and_then(|signal| signal.get("workflow_sequence"))
4300                    .and_then(value_as_u64)
4301            });
4302
4303        signals.push(QuerySignal {
4304            id: signal_id.map(str::to_string).or_else(|| {
4305                matched_export
4306                    .and_then(|signal| signal.get("id"))
4307                    .and_then(Value::as_str)
4308                    .map(str::to_string)
4309            }),
4310            name: name.to_string(),
4311            arguments,
4312            workflow_sequence,
4313        });
4314    }
4315
4316    if signals.is_empty() {
4317        for signal in export_signals {
4318            if signal.get("status").and_then(Value::as_str) == Some("rejected") {
4319                continue;
4320            }
4321            let Some(name) = signal.get("name").and_then(Value::as_str) else {
4322                continue;
4323            };
4324            let codec = signal
4325                .get("payload_codec")
4326                .and_then(Value::as_str)
4327                .unwrap_or(export_codec);
4328            let arguments = decode_query_signal_arguments(signal.get("arguments"), codec)?;
4329            signals.push(QuerySignal {
4330                id: signal.get("id").and_then(Value::as_str).map(str::to_string),
4331                name: name.to_string(),
4332                arguments,
4333                workflow_sequence: signal.get("workflow_sequence").and_then(value_as_u64),
4334            });
4335        }
4336        signals.sort_by_key(|signal| signal.workflow_sequence.unwrap_or(u64::MAX));
4337    }
4338
4339    Ok(signals)
4340}
4341
4342fn decode_query_signal_arguments(raw: Option<&Value>, codec: &str) -> Result<Vec<Value>> {
4343    let decoded = match raw.filter(|value| !value.is_null()) {
4344        Some(value) => decode_wire_value(value, codec)?,
4345        None => Value::Array(Vec::new()),
4346    };
4347    let Value::Array(arguments) = normalize_arguments(decoded) else {
4348        unreachable!("normalize_arguments always returns an array");
4349    };
4350    Ok(arguments)
4351}
4352
4353fn value_as_u64(value: &Value) -> Option<u64> {
4354    value
4355        .as_u64()
4356        .or_else(|| value.as_str().and_then(|value| value.parse().ok()))
4357}
4358
4359fn resume_signal_matches_event(
4360    resume_signal: &ResumeSignal,
4361    event: &HistoryEvent,
4362    signal_name: &str,
4363) -> bool {
4364    if resume_signal.signal_name != signal_name {
4365        return false;
4366    }
4367
4368    match (
4369        resume_signal.signal_id.as_deref(),
4370        event.payload.get("signal_id").and_then(Value::as_str),
4371    ) {
4372        (Some(resume_id), Some(event_id)) => resume_id == event_id,
4373        _ => true,
4374    }
4375}
4376
4377#[cfg(test)]
4378mod tests {
4379    use super::*;
4380    use std::{
4381        io::{Read, Write},
4382        net::{SocketAddr, TcpListener, TcpStream},
4383        thread,
4384    };
4385
4386    #[derive(Clone, Debug, Default, PartialEq)]
4387    struct ReplayCounterState {
4388        loaded: Option<String>,
4389        count: i64,
4390        finished: bool,
4391    }
4392
4393    fn replay_counter_worker() -> Worker {
4394        let client = Client::new("http://127.0.0.1:8080").expect("client");
4395        let mut worker = Worker::new(client, "rust-workers");
4396        worker.register_replayed_workflow(
4397            "replay-counter",
4398            ReplayCounterState::default,
4399            |ctx, _input, state| async move {
4400                let loaded = ctx.activity("load-counter", json!([])).await?;
4401                state.update(|current| {
4402                    current.loaded = loaded.as_str().map(str::to_string);
4403                })?;
4404                for _ in 0..2 {
4405                    let signal = ctx.wait_signal("increment").await?;
4406                    let amount = signal.first().and_then(Value::as_i64).unwrap_or_default();
4407                    state.update(|current| current.count += amount)?;
4408                }
4409                state.update(|current| current.finished = true)?;
4410                state.read(|current| Ok(json!(current.count)))?
4411            },
4412        );
4413        worker.register_replayed_query::<ReplayCounterState, _, _>(
4414            "replay-counter",
4415            "current",
4416            |_ctx, state, _args| async move {
4417                Ok(json!({
4418                    "loaded": state.loaded,
4419                    "count": state.count,
4420                    "finished": state.finished,
4421                }))
4422            },
4423        );
4424        worker.register_replayed_query::<ReplayCounterState, _, _>(
4425            "replay-counter",
4426            "detached-mutation",
4427            |_ctx, state, _args| async move {
4428                let mut detached = (*state).clone();
4429                detached.count = 999;
4430                Ok(json!(detached.count))
4431            },
4432        );
4433        worker.register_replayed_query::<ReplayCounterState, _, _>(
4434            "replay-counter",
4435            "failed-mutation",
4436            |_ctx, state, _args| async move {
4437                let mut detached = (*state).clone();
4438                detached.count = 999;
4439                Err(Error::WorkerLoop("query refused".to_string()))
4440            },
4441        );
4442        worker
4443    }
4444
4445    fn replay_counter_query(
4446        query_name: &str,
4447        history_events: Value,
4448        run_status: &str,
4449    ) -> QueryTask {
4450        serde_json::from_value(json!({
4451            "query_task_id": format!("query-{query_name}"),
4452            "workflow_type": "replay-counter",
4453            "query_name": query_name,
4454            "payload_codec": "json",
4455            "workflow_arguments": {"codec": "json", "blob": "[]"},
4456            "query_arguments": {"codec": "json", "blob": "[]"},
4457            "history_events": history_events,
4458            "run_status": run_status,
4459        }))
4460        .expect("query task")
4461    }
4462
4463    fn workflow_context(history: Vec<HistoryEvent>) -> WorkflowContext {
4464        WorkflowContext {
4465            state: Arc::new(Mutex::new(
4466                WorkflowState::new_with_identity(
4467                    history,
4468                    None,
4469                    None,
4470                    "rust-workers".to_string(),
4471                    JSON_CODEC.to_string(),
4472                    None,
4473                )
4474                .expect("valid workflow history"),
4475            )),
4476        }
4477    }
4478
4479    fn history_event(event_type: &str, payload: Value) -> HistoryEvent {
4480        HistoryEvent {
4481            event_type: event_type.to_string(),
4482            payload,
4483            raw: HashMap::new(),
4484        }
4485    }
4486
4487    #[test]
4488    fn avro_generic_wrapper_round_trips_json_values() {
4489        let value = json!({"greeting": "hello", "count": 3, "ok": true});
4490        let envelope = PayloadEnvelope::avro(&value).expect("encode");
4491        assert_eq!(envelope.codec, DEFAULT_CODEC);
4492        assert_eq!(decode_payload::<Value>(&envelope).expect("decode"), value);
4493    }
4494
4495    #[test]
4496    fn json_codec_remains_plain_json() {
4497        let value = json!({"greeting": "hello", "count": 3, "ok": true});
4498        let envelope = PayloadEnvelope::json(&value).expect("encode");
4499
4500        assert_eq!(envelope.codec, JSON_CODEC);
4501        assert_eq!(envelope.blob, serde_json::to_string(&value).expect("json"));
4502        assert_eq!(decode_payload::<Value>(&envelope).expect("decode"), value);
4503    }
4504
4505    #[test]
4506    fn typed_avro_payload_without_schema_context_keeps_diagnostic() {
4507        let envelope = PayloadEnvelope {
4508            codec: DEFAULT_CODEC.to_string(),
4509            blob: BASE64.encode([0x01]),
4510        };
4511
4512        let error = decode_payload::<Value>(&envelope).expect_err("typed payload must fail");
4513        assert_eq!(
4514            error.to_string(),
4515            "codec error: typed avro payloads require a schema context; v1 supports the generic wrapper"
4516        );
4517    }
4518
4519    #[test]
4520    fn workflow_context_schedules_activity_until_completion_is_in_history() {
4521        let ctx = WorkflowContext {
4522            state: Arc::new(Mutex::new(
4523                WorkflowState::new_with_identity(
4524                    Vec::new(),
4525                    Some("wf-parent".to_string()),
4526                    Some("run-parent".to_string()),
4527                    "rust-workers".to_string(),
4528                    DEFAULT_CODEC.to_string(),
4529                    None,
4530                )
4531                .expect("workflow state"),
4532            )),
4533        };
4534
4535        let mut call = Box::pin(ctx.activity("hello.activity", json!(["Ada"])));
4536        let mut task_context = TaskContext::from_waker(noop_waker_ref());
4537        assert!(matches!(
4538            call.as_mut().poll(&mut task_context),
4539            Poll::Pending
4540        ));
4541
4542        let commands = ctx.take_commands().expect("commands");
4543        assert_eq!(commands[0]["type"], "schedule_activity");
4544        assert_eq!(commands[0]["activity_type"], "hello.activity");
4545    }
4546
4547    #[test]
4548    fn workflow_sleep_emits_one_durable_timer_and_rounds_up() {
4549        let ctx = workflow_context(Vec::new());
4550        let mut sleep = Box::pin(ctx.sleep(Duration::from_millis(1_001)));
4551        let mut task_context = TaskContext::from_waker(noop_waker_ref());
4552
4553        assert!(matches!(
4554            sleep.as_mut().poll(&mut task_context),
4555            Poll::Pending
4556        ));
4557        assert!(matches!(
4558            sleep.as_mut().poll(&mut task_context),
4559            Poll::Pending
4560        ));
4561
4562        let commands = ctx.take_commands().expect("timer command");
4563        assert_eq!(
4564            commands,
4565            vec![json!({
4566                "type": "start_timer",
4567                "delay_seconds": 2,
4568            })]
4569        );
4570    }
4571
4572    #[test]
4573    fn workflow_sleep_replays_matching_schedule_and_fire_without_a_command() {
4574        let history = vec![
4575            history_event(
4576                "TimerScheduled",
4577                json!({
4578                    "sequence": 1,
4579                    "timer_id": "timer-1",
4580                    "delay_seconds": 5,
4581                    "fire_at": "2026-07-11T12:00:05Z",
4582                }),
4583            ),
4584            history_event(
4585                "TimerFired",
4586                json!({
4587                    "sequence": 1,
4588                    "timer_id": "timer-1",
4589                    "delay_seconds": 5,
4590                    "fire_at": "2026-07-11T12:00:05Z",
4591                    "fired_at": "2026-07-11T12:00:05Z",
4592                }),
4593            ),
4594        ];
4595
4596        for _restart in 0..2 {
4597            let ctx = workflow_context(history.clone());
4598            let mut sleep = Box::pin(ctx.sleep(Duration::from_secs(5)));
4599            let mut task_context = TaskContext::from_waker(noop_waker_ref());
4600            assert!(matches!(
4601                sleep.as_mut().poll(&mut task_context),
4602                Poll::Ready(Ok(()))
4603            ));
4604            assert!(ctx.take_commands().expect("commands").is_empty());
4605            ctx.ensure_history_consumed().expect("history consumed");
4606        }
4607    }
4608
4609    #[test]
4610    fn workflow_sleep_rejects_changed_delay_during_replay() {
4611        let ctx = workflow_context(vec![
4612            history_event(
4613                "TimerScheduled",
4614                json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4615            ),
4616            history_event(
4617                "TimerFired",
4618                json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4619            ),
4620        ]);
4621        let mut sleep = Box::pin(ctx.sleep(Duration::from_secs(500)));
4622        let mut task_context = TaskContext::from_waker(noop_waker_ref());
4623
4624        let Poll::Ready(Err(Error::NonDeterministicReplay(failure))) =
4625            sleep.as_mut().poll(&mut task_context)
4626        else {
4627            panic!("changed timer delay must be rejected");
4628        };
4629        assert_eq!(failure.reason, "timer_delay_mismatch");
4630        assert_eq!(failure.sequence, Some(1));
4631    }
4632
4633    #[test]
4634    fn workflow_history_rejects_unpaired_or_mismatched_timer_events() {
4635        let lone_fire = WorkflowState::new(
4636            vec![history_event(
4637                "TimerFired",
4638                json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4639            )],
4640            "rust-workers".to_string(),
4641            JSON_CODEC.to_string(),
4642            None,
4643        )
4644        .expect_err("TimerFired requires TimerScheduled");
4645        assert!(matches!(
4646            lone_fire,
4647            Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4648                if reason == "timer_schedule_missing_or_duplicate"
4649        ));
4650
4651        let wrong_identity = WorkflowState::new(
4652            vec![
4653                history_event(
4654                    "TimerScheduled",
4655                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4656                ),
4657                history_event(
4658                    "TimerFired",
4659                    json!({"sequence": 1, "timer_id": "timer-2", "delay_seconds": 5}),
4660                ),
4661            ],
4662            "rust-workers".to_string(),
4663            JSON_CODEC.to_string(),
4664            None,
4665        )
4666        .expect_err("fire must match scheduled timer identity");
4667        assert!(matches!(
4668            wrong_identity,
4669            Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4670                if reason == "timer_identity_mismatch"
4671        ));
4672
4673        let duplicate_fire = WorkflowState::new(
4674            vec![
4675                history_event(
4676                    "TimerScheduled",
4677                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4678                ),
4679                history_event(
4680                    "TimerFired",
4681                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4682                ),
4683                history_event(
4684                    "TimerFired",
4685                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4686                ),
4687            ],
4688            "rust-workers".to_string(),
4689            JSON_CODEC.to_string(),
4690            None,
4691        )
4692        .expect_err("a durable timer cannot fire twice");
4693        assert!(matches!(
4694            duplicate_fire,
4695            Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4696                if reason == "duplicate_timer_fire"
4697        ));
4698
4699        let wrong_fired_delay = WorkflowState::new(
4700            vec![
4701                history_event(
4702                    "TimerScheduled",
4703                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4704                ),
4705                history_event(
4706                    "TimerFired",
4707                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 6}),
4708                ),
4709            ],
4710            "rust-workers".to_string(),
4711            JSON_CODEC.to_string(),
4712            None,
4713        )
4714        .expect_err("timer schedule and fire delays must agree");
4715        assert!(matches!(
4716            wrong_fired_delay,
4717            Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4718                if reason == "timer_history_delay_mismatch"
4719        ));
4720    }
4721
4722    #[test]
4723    fn replay_rejects_activity_moved_before_recorded_timer() {
4724        let ctx = workflow_context(vec![
4725            history_event(
4726                "TimerScheduled",
4727                json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4728            ),
4729            history_event(
4730                "TimerFired",
4731                json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4732            ),
4733            history_event(
4734                "ActivityCompleted",
4735                json!({
4736                    "sequence": 2,
4737                    "activity_type": "after-timer",
4738                    "payload_codec": "json",
4739                    "result": {"codec": "json", "blob": "\"done\""},
4740                }),
4741            ),
4742        ]);
4743        let mut activity = Box::pin(ctx.activity("after-timer", json!([])));
4744        let mut task_context = TaskContext::from_waker(noop_waker_ref());
4745
4746        let Poll::Ready(Err(Error::NonDeterministicReplay(failure))) =
4747            activity.as_mut().poll(&mut task_context)
4748        else {
4749            panic!("reordered durable command must be rejected");
4750        };
4751        assert_eq!(failure.reason, "recorded_command_mismatch");
4752        assert_eq!(failure.sequence, Some(1));
4753        assert_eq!(failure.expected.as_deref(), Some("timer"));
4754        assert_eq!(failure.actual.as_deref(), Some("activity:after-timer"));
4755    }
4756
4757    #[test]
4758    fn replay_orders_signal_waits_and_timers_in_one_command_stream() {
4759        let signal_then_timer = vec![
4760            history_event(
4761                "ConditionWaitOpened",
4762                json!({"sequence": 1, "condition_key": "signal:go"}),
4763            ),
4764            history_event(
4765                "SignalReceived",
4766                json!({
4767                    "signal_name": "go",
4768                    "arguments": ["now"],
4769                }),
4770            ),
4771            history_event(
4772                "TimerScheduled",
4773                json!({"sequence": 2, "timer_id": "timer-2", "delay_seconds": 5}),
4774            ),
4775            history_event(
4776                "TimerFired",
4777                json!({"sequence": 2, "timer_id": "timer-2", "delay_seconds": 5}),
4778            ),
4779        ];
4780
4781        let ctx = workflow_context(signal_then_timer.clone());
4782        let mut signal = Box::pin(ctx.wait_signal("go"));
4783        let mut task_context = TaskContext::from_waker(noop_waker_ref());
4784        assert!(matches!(
4785            signal.as_mut().poll(&mut task_context),
4786            Poll::Ready(Ok(arguments)) if arguments == vec![json!("now")]
4787        ));
4788        let mut timer = Box::pin(ctx.sleep(Duration::from_secs(5)));
4789        assert!(matches!(
4790            timer.as_mut().poll(&mut task_context),
4791            Poll::Ready(Ok(()))
4792        ));
4793        ctx.ensure_history_consumed()
4794            .expect("signal and timer history consumed in order");
4795
4796        let reordered = workflow_context(signal_then_timer);
4797        let mut timer_first = Box::pin(reordered.sleep(Duration::from_secs(5)));
4798        let Poll::Ready(Err(Error::NonDeterministicReplay(failure))) =
4799            timer_first.as_mut().poll(&mut task_context)
4800        else {
4801            panic!("timer cannot consume signal-wait-first history");
4802        };
4803        assert_eq!(failure.reason, "recorded_command_mismatch");
4804        assert_eq!(failure.sequence, Some(1));
4805        assert_eq!(failure.expected.as_deref(), Some("signal wait"));
4806
4807        let timer_then_signal = vec![
4808            history_event(
4809                "TimerScheduled",
4810                json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4811            ),
4812            history_event(
4813                "TimerFired",
4814                json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4815            ),
4816            history_event(
4817                "ConditionWaitOpened",
4818                json!({"sequence": 2, "condition_key": "signal:go"}),
4819            ),
4820            history_event(
4821                "SignalReceived",
4822                json!({"signal_name": "go", "arguments": []}),
4823            ),
4824        ];
4825        let reordered = workflow_context(timer_then_signal);
4826        let mut signal_first = Box::pin(reordered.wait_signal("go"));
4827        let Poll::Ready(Err(Error::NonDeterministicReplay(failure))) =
4828            signal_first.as_mut().poll(&mut task_context)
4829        else {
4830            panic!("signal wait cannot consume timer-first history");
4831        };
4832        assert_eq!(failure.reason, "recorded_command_mismatch");
4833        assert_eq!(failure.sequence, Some(1));
4834        assert_eq!(failure.expected.as_deref(), Some("timer"));
4835    }
4836
4837    #[test]
4838    fn workflow_history_rejects_duplicate_or_colliding_command_sequences() {
4839        let duplicate_timer = WorkflowState::new(
4840            vec![
4841                history_event(
4842                    "TimerScheduled",
4843                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4844                ),
4845                history_event(
4846                    "TimerScheduled",
4847                    json!({"sequence": 1, "timer_id": "timer-2", "delay_seconds": 5}),
4848                ),
4849            ],
4850            "rust-workers".to_string(),
4851            JSON_CODEC.to_string(),
4852            None,
4853        )
4854        .expect_err("one workflow sequence cannot schedule two timers");
4855        assert!(matches!(
4856            duplicate_timer,
4857            Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4858                if reason == "timer_schedule_missing_or_duplicate"
4859        ));
4860
4861        let colliding_kinds = WorkflowState::new(
4862            vec![
4863                history_event(
4864                    "TimerScheduled",
4865                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4866                ),
4867                history_event(
4868                    "ActivityCompleted",
4869                    json!({"sequence": 1, "activity_type": "same-sequence"}),
4870                ),
4871            ],
4872            "rust-workers".to_string(),
4873            JSON_CODEC.to_string(),
4874            None,
4875        )
4876        .expect_err("one workflow sequence cannot identify two command kinds");
4877        assert!(matches!(
4878            colliding_kinds,
4879            Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4880                if reason == "durable_command_sequence_collision"
4881        ));
4882
4883        let duplicate_signal_wait = WorkflowState::new(
4884            vec![
4885                history_event(
4886                    "SignalWaitOpened",
4887                    json!({"sequence": 1, "signal_name": "go"}),
4888                ),
4889                history_event(
4890                    "SignalWaitOpened",
4891                    json!({"sequence": 1, "signal_name": "go"}),
4892                ),
4893            ],
4894            "rust-workers".to_string(),
4895            JSON_CODEC.to_string(),
4896            None,
4897        )
4898        .expect_err("one workflow sequence cannot open two signal waits");
4899        assert!(matches!(
4900            duplicate_signal_wait,
4901            Error::NonDeterministicReplay(ReplayFailure { ref reason, .. })
4902                if reason == "duplicate_signal_wait_open"
4903        ));
4904    }
4905
4906    #[test]
4907    fn workflow_sleep_rejects_unrepresentable_rounded_duration() {
4908        let ctx = workflow_context(Vec::new());
4909        let mut sleep = Box::pin(ctx.start_timer(Duration::new(u64::MAX, 1)));
4910        let mut task_context = TaskContext::from_waker(noop_waker_ref());
4911        assert!(matches!(
4912            sleep.as_mut().poll(&mut task_context),
4913            Poll::Ready(Err(Error::TimerDurationOverflow))
4914        ));
4915        assert!(ctx.take_commands().expect("commands").is_empty());
4916    }
4917
4918    #[test]
4919    fn workflow_task_replay_completes_without_rescheduling_recorded_commands() {
4920        let client = Client::new("http://127.0.0.1:8080").expect("client");
4921        let mut worker = Worker::new(client, "rust-workers");
4922        worker.register_workflow("rust.timer", |ctx, _input| async move {
4923            ctx.sleep(Duration::from_secs(5)).await?;
4924            ctx.activity("after-timer", json!([])).await
4925        });
4926
4927        let task = |history_events| WorkflowTask {
4928            task_id: "wft-rust-timer-1".to_string(),
4929            workflow_id: Some("wf-rust-timer".to_string()),
4930            run_id: Some("run-rust-timer".to_string()),
4931            workflow_type: "rust.timer".to_string(),
4932            payload_codec: JSON_CODEC.to_string(),
4933            arguments: Some(json!({"codec": "json", "blob": "[]"})),
4934            history_events,
4935            total_history_events: None,
4936            next_history_page_token: None,
4937            workflow_task_attempt: 1,
4938            workflow_signal_id: None,
4939            signal_name: None,
4940            signal_arguments: None,
4941            lease_owner: Some("rust-worker".to_string()),
4942        };
4943
4944        let initial = worker
4945            .execute_workflow_task(task(Vec::new()))
4946            .expect("initial timer task");
4947        assert_eq!(
4948            initial,
4949            vec![json!({"type": "start_timer", "delay_seconds": 5})]
4950        );
4951
4952        let replayed = worker
4953            .execute_workflow_task(task(vec![
4954                history_event(
4955                    "TimerScheduled",
4956                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4957                ),
4958                history_event(
4959                    "TimerFired",
4960                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
4961                ),
4962                history_event(
4963                    "ActivityCompleted",
4964                    json!({
4965                        "sequence": 2,
4966                        "activity_type": "after-timer",
4967                        "payload_codec": "json",
4968                        "result": {"codec": "json", "blob": "\"done\""},
4969                    }),
4970                ),
4971            ]))
4972            .expect("replayed workflow task");
4973        assert_eq!(replayed.len(), 1);
4974        assert_eq!(replayed[0]["type"], "complete_workflow");
4975        assert_eq!(
4976            decode_wire_value(&replayed[0]["result"], JSON_CODEC).expect("result"),
4977            json!("done")
4978        );
4979    }
4980
4981    #[test]
4982    fn workflow_task_replay_keeps_recorded_unfired_timer_pending_without_rescheduling() {
4983        let client = Client::new("http://127.0.0.1:8080").expect("client");
4984        let mut worker = Worker::new(client, "rust-workers");
4985        worker.register_workflow("rust.timer.pending", |ctx, _input| async move {
4986            ctx.sleep(Duration::from_secs(5)).await?;
4987            Ok(json!({"status": "timer fired"}))
4988        });
4989
4990        let task = WorkflowTask {
4991            task_id: "wft-rust-timer-pending".to_string(),
4992            workflow_id: Some("wf-rust-timer".to_string()),
4993            run_id: Some("run-rust-timer".to_string()),
4994            workflow_type: "rust.timer.pending".to_string(),
4995            payload_codec: JSON_CODEC.to_string(),
4996            arguments: Some(json!({"codec": "json", "blob": "[]"})),
4997            history_events: vec![history_event(
4998                "TimerScheduled",
4999                json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
5000            )],
5001            total_history_events: Some(1),
5002            next_history_page_token: None,
5003            workflow_task_attempt: 1,
5004            workflow_signal_id: None,
5005            signal_name: None,
5006            signal_arguments: None,
5007            lease_owner: Some("rust-worker".to_string()),
5008        };
5009
5010        for _redelivery_or_restart in 0..2 {
5011            let commands = worker
5012                .execute_workflow_task(task.clone())
5013                .expect("recorded timer remains pending");
5014            assert!(
5015                commands.is_empty(),
5016                "recorded timer must not be rescheduled"
5017            );
5018        }
5019    }
5020
5021    #[test]
5022    fn workflow_task_rejects_recorded_command_removed_from_workflow_code() {
5023        let client = Client::new("http://127.0.0.1:8080").expect("client");
5024        let mut worker = Worker::new(client, "rust-workers");
5025        worker.register_workflow("rust.timer.removed", |_ctx, _input| async move {
5026            Ok(json!({"status": "completed"}))
5027        });
5028        let task = WorkflowTask {
5029            task_id: "wft-rust-timer-removed".to_string(),
5030            workflow_id: Some("wf-rust-timer".to_string()),
5031            run_id: Some("run-rust-timer".to_string()),
5032            workflow_type: "rust.timer.removed".to_string(),
5033            payload_codec: JSON_CODEC.to_string(),
5034            arguments: Some(json!({"codec": "json", "blob": "[]"})),
5035            history_events: vec![
5036                history_event(
5037                    "TimerScheduled",
5038                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
5039                ),
5040                history_event(
5041                    "TimerFired",
5042                    json!({"sequence": 1, "timer_id": "timer-1", "delay_seconds": 5}),
5043                ),
5044            ],
5045            total_history_events: Some(2),
5046            next_history_page_token: None,
5047            workflow_task_attempt: 1,
5048            workflow_signal_id: None,
5049            signal_name: None,
5050            signal_arguments: None,
5051            lease_owner: Some("rust-worker".to_string()),
5052        };
5053
5054        let Error::NonDeterministicReplay(failure) = worker
5055            .execute_workflow_task(task)
5056            .expect_err("removed timer must fail replay")
5057        else {
5058            panic!("expected typed replay failure");
5059        };
5060        assert_eq!(failure.reason, "recorded_commands_unconsumed");
5061        assert_eq!(failure.sequence, Some(1));
5062    }
5063
5064    #[test]
5065    fn workflow_context_emits_explicit_child_workflow_contract() {
5066        let ctx = WorkflowContext {
5067            state: Arc::new(Mutex::new(
5068                WorkflowState::new_with_identity(
5069                    Vec::new(),
5070                    Some("wf-parent".to_string()),
5071                    Some("run-parent".to_string()),
5072                    "parent-workers".to_string(),
5073                    JSON_CODEC.to_string(),
5074                    None,
5075                )
5076                .expect("workflow state"),
5077            )),
5078        };
5079        let options = ChildWorkflowOptions::new("python-workers")
5080            .parent_close_policy(ParentClosePolicy::RequestCancel)
5081            .retry_policy(ChildWorkflowRetryPolicy {
5082                max_attempts: Some(3),
5083                backoff_seconds: vec![1, 5],
5084                non_retryable_error_types: vec!["ValidationError".to_string()],
5085            })
5086            .execution_timeout_seconds(600)
5087            .run_timeout_seconds(120);
5088        let mut call = Box::pin(ctx.start_child_workflow(
5089            "python.fulfil-order",
5090            options,
5091            json!([{"order_id": "order-42"}]),
5092        ));
5093        let mut task_context = TaskContext::from_waker(noop_waker_ref());
5094
5095        assert!(matches!(
5096            call.as_mut().poll(&mut task_context),
5097            Poll::Pending
5098        ));
5099        let commands = ctx.take_commands().expect("commands");
5100        assert_eq!(commands.len(), 1);
5101        let command = &commands[0];
5102        assert_eq!(command["type"], "start_child_workflow");
5103        assert_eq!(command["workflow_type"], "python.fulfil-order");
5104        assert_eq!(command["queue"], "python-workers");
5105        assert_eq!(command["parent_close_policy"], "request_cancel");
5106        assert_eq!(command["retry_policy"]["max_attempts"], 3);
5107        assert_eq!(command["execution_timeout_seconds"], 600);
5108        assert_eq!(command["run_timeout_seconds"], 120);
5109        assert_eq!(
5110            decode_wire_value(&command["arguments"], JSON_CODEC).expect("child args"),
5111            json!([{"order_id": "order-42"}])
5112        );
5113    }
5114
5115    fn child_parent_worker() -> Worker {
5116        let client = Client::new("http://127.0.0.1:8080").expect("client");
5117        let mut worker = Worker::new(client, "rust-parent-workers");
5118        worker.register_workflow("rust.parent", |ctx, _input| async move {
5119            let child = ctx
5120                .start_child_workflow(
5121                    "python.child",
5122                    ChildWorkflowOptions::new("python-child-workers")
5123                        .parent_close_policy(ParentClosePolicy::Terminate),
5124                    json!([{"codec_probe": [1, true, "rust"]}]),
5125                )
5126                .await?;
5127            Ok(json!({
5128                "parent_workflow_id": child.parent.workflow_id,
5129                "parent_run_id": child.parent.run_id,
5130                "child_workflow_id": child.child.workflow_id,
5131                "child_run_id": child.child.run_id,
5132                "child_workflow_type": child.child_workflow_type,
5133                "result": child.result,
5134            }))
5135        });
5136        worker
5137    }
5138
5139    fn child_parent_task(event_type: &str, payload: Value) -> WorkflowTask {
5140        WorkflowTask {
5141            task_id: "wft-child-parent".to_string(),
5142            workflow_id: Some("wf-parent".to_string()),
5143            run_id: Some("run-parent".to_string()),
5144            workflow_type: "rust.parent".to_string(),
5145            payload_codec: JSON_CODEC.to_string(),
5146            arguments: Some(encode_value_envelope(&json!([]), JSON_CODEC).expect("input")),
5147            history_events: vec![
5148                HistoryEvent {
5149                    event_type: "ChildWorkflowScheduled".to_string(),
5150                    payload: json!({
5151                        "sequence": 1,
5152                        "child_call_id": "call-child",
5153                        "child_workflow_instance_id": "wf-child",
5154                        "child_workflow_run_id": "run-child",
5155                        "child_workflow_type": "python.child",
5156                    }),
5157                    raw: HashMap::new(),
5158                },
5159                HistoryEvent {
5160                    event_type: event_type.to_string(),
5161                    payload,
5162                    raw: HashMap::new(),
5163                },
5164            ],
5165            total_history_events: Some(2),
5166            next_history_page_token: None,
5167            workflow_task_attempt: 1,
5168            workflow_signal_id: None,
5169            signal_name: None,
5170            signal_arguments: None,
5171            lease_owner: Some("rust-worker".to_string()),
5172        }
5173    }
5174
5175    #[test]
5176    fn committed_child_result_replays_without_starting_a_duplicate() {
5177        let worker = child_parent_worker();
5178        let task = child_parent_task(
5179            "ChildRunCompleted",
5180            json!({
5181                "sequence": 1,
5182                "child_call_id": "call-child",
5183                "child_workflow_instance_id": "wf-child",
5184                "child_workflow_run_id": "run-child",
5185                "child_workflow_type": "python.child",
5186                "payload_codec": "json",
5187                "result": {"codec": "json", "blob": "{\"from\":\"python\",\"ok\":true}"},
5188            }),
5189        );
5190
5191        for _restart in 0..2 {
5192            let commands = worker
5193                .execute_workflow_task(task.clone())
5194                .expect("replayed parent task");
5195            assert_eq!(commands.len(), 1);
5196            assert_eq!(commands[0]["type"], "complete_workflow");
5197            assert!(!commands
5198                .iter()
5199                .any(|command| command["type"] == "start_child_workflow"));
5200            let output =
5201                decode_wire_value(&commands[0]["result"], JSON_CODEC).expect("parent output");
5202            assert_eq!(output["parent_workflow_id"], "wf-parent");
5203            assert_eq!(output["parent_run_id"], "run-parent");
5204            assert_eq!(output["child_workflow_id"], "wf-child");
5205            assert_eq!(output["child_run_id"], "run-child");
5206            assert_eq!(output["result"], json!({"from": "python", "ok": true}));
5207        }
5208    }
5209
5210    #[test]
5211    fn pending_child_replays_after_restart_without_starting_a_duplicate() {
5212        let worker = child_parent_worker();
5213        let mut task = child_parent_task("unused", Value::Null);
5214        task.history_events.truncate(1);
5215        task.total_history_events = Some(1);
5216
5217        for _redelivery_or_restart in 0..2 {
5218            let commands = worker
5219                .execute_workflow_task(task.clone())
5220                .expect("recorded child remains pending");
5221            assert!(
5222                commands.is_empty(),
5223                "recorded pending child must not be started again"
5224            );
5225        }
5226    }
5227
5228    #[test]
5229    fn child_cancellation_becomes_stable_parent_failure_command() {
5230        let worker = child_parent_worker();
5231        let task = child_parent_task(
5232            "ChildRunCancelled",
5233            json!({
5234                "sequence": 1,
5235                "child_workflow_instance_id": "wf-child",
5236                "child_workflow_run_id": "run-child",
5237                "child_workflow_type": "python.child",
5238                "failure_id": "failure-child",
5239                "failure_category": "cancelled",
5240                "message": "cancelled by parent-close policy",
5241            }),
5242        );
5243
5244        let commands = worker
5245            .execute_workflow_task(task)
5246            .expect("parent settlement");
5247        assert_eq!(commands.len(), 1);
5248        assert_eq!(commands[0]["type"], "fail_workflow");
5249        assert_eq!(commands[0]["exception_type"], "ChildWorkflowCancelled");
5250        assert_eq!(
5251            commands[0]["exception"]["properties"]["reason"],
5252            "cancelled"
5253        );
5254        assert_eq!(
5255            commands[0]["exception"]["properties"]["child_workflow_run_id"],
5256            "run-child"
5257        );
5258    }
5259
5260    #[test]
5261    fn workflow_can_handle_typed_child_failure() {
5262        let client = Client::new("http://127.0.0.1:8080").expect("client");
5263        let mut worker = Worker::new(client, "rust-parent-workers");
5264        worker.register_workflow("rust.handled-parent", |ctx, _input| async move {
5265            match ctx
5266                .start_child_workflow(
5267                    "python.child",
5268                    ChildWorkflowOptions::new("python-child-workers"),
5269                    json!([]),
5270                )
5271                .await
5272            {
5273                Err(Error::ChildWorkflowFailed(failure)) => Ok(json!({
5274                    "reason": failure.reason,
5275                    "failure_id": failure.failure_id,
5276                    "exception_class": failure.exception_class,
5277                    "child_run_id": failure.child_workflow_run_id,
5278                })),
5279                Err(error) => Err(error),
5280                Ok(_) => Err(Error::WorkerLoop(
5281                    "child unexpectedly succeeded".to_string(),
5282                )),
5283            }
5284        });
5285        let mut task = child_parent_task(
5286            "ChildRunFailed",
5287            json!({
5288                "sequence": 1,
5289                "child_workflow_instance_id": "wf-child",
5290                "child_workflow_run_id": "run-child",
5291                "child_workflow_type": "python.child",
5292                "failure_id": "failure-child",
5293                "failure_category": "child_workflow",
5294                "message": "payment rejected",
5295                "exception": {
5296                    "type": "PaymentRejected",
5297                    "class": "payments.PaymentRejected",
5298                    "message": "payment rejected"
5299                }
5300            }),
5301        );
5302        task.workflow_type = "rust.handled-parent".to_string();
5303
5304        let commands = worker.execute_workflow_task(task).expect("handled failure");
5305        assert_eq!(commands[0]["type"], "complete_workflow");
5306        let output = decode_wire_value(&commands[0]["result"], JSON_CODEC).expect("parent output");
5307        assert_eq!(output["reason"], "child_workflow");
5308        assert_eq!(output["failure_id"], "failure-child");
5309        assert_eq!(output["exception_class"], "payments.PaymentRejected");
5310        assert_eq!(output["child_run_id"], "run-child");
5311    }
5312
5313    #[test]
5314    fn rust_hello_world_uses_signal_arguments_from_resume_payload() {
5315        let client = Client::new("http://127.0.0.1:8080").expect("client");
5316        let mut worker = Worker::new(client, "rust-workers");
5317
5318        worker.register_workflow("rust.hello_workflow", |ctx, _input| async move {
5319            let signal = ctx.wait_signal("start").await?;
5320            let name = signal
5321                .first()
5322                .and_then(|value| value.as_str())
5323                .unwrap_or("world");
5324            let greeting = ctx.activity("rust.hello_activity", json!([name])).await?;
5325            Ok(json!({
5326                "greeting": greeting,
5327                "language": "rust"
5328            }))
5329        });
5330
5331        let signal_arguments =
5332            encode_value_envelope(&json!(["Rust"]), DEFAULT_CODEC).expect("signal arguments");
5333        let task = WorkflowTask {
5334            task_id: "wft-rust-signal-1".to_string(),
5335            workflow_id: Some("wf-rust-hello".to_string()),
5336            run_id: Some("run-rust-hello".to_string()),
5337            workflow_type: "rust.hello_workflow".to_string(),
5338            payload_codec: DEFAULT_CODEC.to_string(),
5339            arguments: Some(encode_value_envelope(&json!([]), DEFAULT_CODEC).expect("input")),
5340            history_events: vec![HistoryEvent {
5341                event_type: "SignalReceived".to_string(),
5342                payload: json!({
5343                    "signal_id": "sig-rust-1",
5344                    "signal_name": "start"
5345                }),
5346                raw: HashMap::new(),
5347            }],
5348            total_history_events: Some(1),
5349            next_history_page_token: None,
5350            workflow_task_attempt: 1,
5351            workflow_signal_id: Some("sig-rust-1".to_string()),
5352            signal_name: Some("start".to_string()),
5353            signal_arguments: Some(signal_arguments),
5354            lease_owner: Some("rust-worker".to_string()),
5355        };
5356
5357        let commands = worker.execute_workflow_task(task).expect("workflow task");
5358
5359        assert_eq!(commands.len(), 1);
5360        assert_eq!(commands[0]["type"], "schedule_activity");
5361        assert_eq!(commands[0]["activity_type"], "rust.hello_activity");
5362        assert_eq!(
5363            decode_wire_value(&commands[0]["arguments"], DEFAULT_CODEC).expect("activity args"),
5364            json!(["Rust"])
5365        );
5366    }
5367
5368    #[test]
5369    fn workflow_task_appends_paginated_history_events() {
5370        let mut task = WorkflowTask {
5371            task_id: "wft-rust-pages-1".to_string(),
5372            workflow_id: Some("wf-rust-pages".to_string()),
5373            run_id: Some("run-rust-pages".to_string()),
5374            workflow_type: "rust.hello_workflow".to_string(),
5375            payload_codec: DEFAULT_CODEC.to_string(),
5376            arguments: Some(encode_value_envelope(&json!([]), DEFAULT_CODEC).expect("input")),
5377            history_events: vec![HistoryEvent {
5378                event_type: "WorkflowStarted".to_string(),
5379                payload: json!({}),
5380                raw: HashMap::new(),
5381            }],
5382            total_history_events: Some(3),
5383            next_history_page_token: Some("MQ==".to_string()),
5384            workflow_task_attempt: 1,
5385            workflow_signal_id: None,
5386            signal_name: None,
5387            signal_arguments: None,
5388            lease_owner: Some("rust-worker".to_string()),
5389        };
5390
5391        task.append_history_page(WorkflowTaskHistoryPage {
5392            history_events: vec![
5393                HistoryEvent {
5394                    event_type: "SignalReceived".to_string(),
5395                    payload: json!({
5396                        "signal_id": "sig-rust-1",
5397                        "signal_name": "start",
5398                        "arguments": encode_value_envelope(&json!(["Rust"]), DEFAULT_CODEC)
5399                            .expect("signal arguments")
5400                    }),
5401                    raw: HashMap::new(),
5402                },
5403                HistoryEvent {
5404                    event_type: "MarkerRecorded".to_string(),
5405                    payload: json!({"sequence": 3}),
5406                    raw: HashMap::new(),
5407                },
5408            ],
5409            total_history_events: Some(3),
5410            next_history_page_token: None,
5411        });
5412
5413        assert_eq!(task.history_events.len(), 3);
5414        assert_eq!(task.total_history_events, Some(3));
5415        assert_eq!(task.next_history_page_token, None);
5416
5417        let signals =
5418            signal_values(&task.history_events, "start", DEFAULT_CODEC, None).expect("signals");
5419        assert_eq!(signals, vec![vec![json!("Rust")]]);
5420    }
5421
5422    #[tokio::test]
5423    async fn query_handler_reads_ordered_cross_codec_signals_without_commands() {
5424        let client = Client::new("http://127.0.0.1:8080").expect("client");
5425        let mut worker = Worker::new(client, "rust-workers");
5426        worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5427        worker.register_query("counter", "current", |ctx, _args| async move {
5428            let mut count = 0_i64;
5429            for signal in ctx.signal_events() {
5430                let value = signal
5431                    .arguments
5432                    .first()
5433                    .and_then(Value::as_i64)
5434                    .unwrap_or_default();
5435                match signal.name.as_str() {
5436                    "increment" => count += value,
5437                    "set" => count = value,
5438                    _ => {}
5439                }
5440            }
5441            Ok(json!(count))
5442        });
5443
5444        let task = QueryTask {
5445            query_task_id: "query-rust-counter".to_string(),
5446            query_task_attempt: 1,
5447            lease_owner: Some("rust-worker".to_string()),
5448            workflow_id: Some("counter-1".to_string()),
5449            run_id: Some("run-counter-1".to_string()),
5450            workflow_type: "counter".to_string(),
5451            query_name: "current".to_string(),
5452            payload_codec: DEFAULT_CODEC.to_string(),
5453            workflow_arguments: Some(
5454                encode_value_envelope(&json!([]), DEFAULT_CODEC).expect("workflow input"),
5455            ),
5456            query_arguments: Some(
5457                encode_value_envelope(&json!([]), DEFAULT_CODEC).expect("query arguments"),
5458            ),
5459            history_events: vec![
5460                HistoryEvent {
5461                    event_type: "SignalReceived".to_string(),
5462                    payload: json!({
5463                        "signal_id": "php-signal-1",
5464                        "signal_name": "increment",
5465                        "workflow_sequence": 1,
5466                        "payload_codec": DEFAULT_CODEC,
5467                        "arguments": encode_value_envelope(&json!([3]), DEFAULT_CODEC).expect("php avro signal")
5468                    }),
5469                    raw: HashMap::new(),
5470                },
5471                HistoryEvent {
5472                    event_type: "SignalReceived".to_string(),
5473                    payload: json!({
5474                        "signal_id": "python-signal-2",
5475                        "signal_name": "increment",
5476                        "workflow_sequence": 2,
5477                        "payload_codec": JSON_CODEC,
5478                        "arguments": encode_value_envelope(&json!([5]), JSON_CODEC).expect("python json signal")
5479                    }),
5480                    raw: HashMap::new(),
5481                },
5482                HistoryEvent {
5483                    event_type: "SignalReceived".to_string(),
5484                    payload: json!({
5485                        "signal_id": "rust-signal-3",
5486                        "signal_name": "set",
5487                        "workflow_sequence": 3,
5488                        "payload_codec": DEFAULT_CODEC,
5489                        "arguments": encode_value_envelope(&json!([0]), DEFAULT_CODEC).expect("rust avro signal")
5490                    }),
5491                    raw: HashMap::new(),
5492                },
5493            ],
5494            history_export: None,
5495            run_status: Some("completed".to_string()),
5496        };
5497
5498        let result = worker.execute_query_task(task).await.expect("query result");
5499        assert_eq!(result, json!(0));
5500    }
5501
5502    #[tokio::test]
5503    async fn replayed_queries_read_running_completed_and_cold_restarted_instance_state() {
5504        let worker = replay_counter_worker();
5505        let running_history = json!([
5506            {
5507                "type": "ActivityCompleted",
5508                "payload": {
5509                    "sequence": 1,
5510                    "activity_type": "load-counter",
5511                    "payload_codec": "json",
5512                    "result": {"codec": "json", "blob": "\"loaded\""}
5513                }
5514            },
5515            {
5516                "type": "SignalReceived",
5517                "payload": {
5518                    "signal_id": "signal-3",
5519                    "signal_name": "increment",
5520                    "payload_codec": "json",
5521                    "arguments": {"codec": "json", "blob": "[3]"}
5522                }
5523            }
5524        ]);
5525
5526        let running = worker
5527            .execute_query_task(replay_counter_query(
5528                "current",
5529                running_history.clone(),
5530                "running",
5531            ))
5532            .await
5533            .expect("running replay query");
5534        assert_eq!(
5535            running,
5536            json!({"loaded": "loaded", "count": 3, "finished": false})
5537        );
5538
5539        let detached = worker
5540            .execute_query_task(replay_counter_query(
5541                "detached-mutation",
5542                running_history.clone(),
5543                "running",
5544            ))
5545            .await
5546            .expect("query mutates only its detached state clone");
5547        assert_eq!(detached, json!(999));
5548        let failed = worker
5549            .execute_query_task(replay_counter_query(
5550                "failed-mutation",
5551                running_history.clone(),
5552                "running",
5553            ))
5554            .await
5555            .expect_err("failed query");
5556        assert_eq!(failed.reason, "query_rejected");
5557        let unchanged = worker
5558            .execute_query_task(replay_counter_query("current", running_history, "running"))
5559            .await
5560            .expect("later query reconstructs unchanged state");
5561        assert_eq!(unchanged, running);
5562
5563        let restarted_worker = replay_counter_worker();
5564        let restarted_task: QueryTask = serde_json::from_value(json!({
5565            "query_task_id": "query-after-restart",
5566            "workflow_id": "counter-1",
5567            "run_id": "run-counter-1",
5568            "workflow_type": "replay-counter",
5569            "query_name": "current",
5570            "payload_codec": "json",
5571            "workflow_arguments": {"codec": "json", "blob": "[]"},
5572            "query_arguments": {"codec": "json", "blob": "[]"},
5573            "history_events": [],
5574            "history_export": {
5575                "payloads": {"codec": "json"},
5576                "history_events": [
5577                    {
5578                        "type": "ActivityCompleted",
5579                        "payload": {
5580                            "sequence": 1,
5581                            "activity_type": "load-counter",
5582                            "payload_codec": "json",
5583                            "result": null
5584                        }
5585                    },
5586                    {
5587                        "type": "SignalReceived",
5588                        "payload": {"signal_id": "signal-3", "signal_name": "increment"}
5589                    },
5590                    {
5591                        "type": "SignalReceived",
5592                        "payload": {"signal_id": "signal-5", "signal_name": "increment"}
5593                    }
5594                ],
5595                "activities": [{
5596                    "sequence": 1,
5597                    "activity_type": "load-counter",
5598                    "payload_codec": "json",
5599                    "result": {"codec": "json", "blob": "\"loaded\""}
5600                }],
5601                "signals": [
5602                    {
5603                        "id": "signal-3",
5604                        "name": "increment",
5605                        "payload_codec": "json",
5606                        "arguments": "[3]"
5607                    },
5608                    {
5609                        "id": "signal-5",
5610                        "name": "increment",
5611                        "payload_codec": "json",
5612                        "arguments": "[5]"
5613                    }
5614                ]
5615            },
5616            "run_status": "completed"
5617        }))
5618        .expect("cold replay query task");
5619        let completed = restarted_worker
5620            .execute_query_task(restarted_task)
5621            .await
5622            .expect("completed cold replay query");
5623        assert_eq!(
5624            completed,
5625            json!({"loaded": "loaded", "count": 8, "finished": true})
5626        );
5627    }
5628
5629    #[tokio::test]
5630    async fn replayed_query_replay_failures_are_machine_readable() {
5631        let worker = replay_counter_worker();
5632        let task = replay_counter_query(
5633            "current",
5634            json!([{
5635                "type": "ActivityCompleted",
5636                "payload": {
5637                    "sequence": 1,
5638                    "payload_codec": "json",
5639                    "result": {"codec": "json", "blob": "{"}
5640                }
5641            }]),
5642            "running",
5643        );
5644        let failure = worker
5645            .execute_query_task(task)
5646            .await
5647            .expect_err("invalid replay history payload");
5648        assert_eq!(failure.reason, "query_workflow_state_unavailable");
5649        assert_eq!(failure.failure_type, "QueryWorkflowStateUnavailable");
5650    }
5651
5652    #[tokio::test]
5653    async fn query_task_restores_compact_history_from_export() {
5654        let client = Client::new("http://127.0.0.1:8080").expect("client");
5655        let mut worker = Worker::new(client, "rust-workers");
5656        worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5657        worker.register_query("counter", "current", |ctx, _args| async move {
5658            Ok(json!(ctx.signals("increment")[0][0]))
5659        });
5660        let task: QueryTask = serde_json::from_value(json!({
5661            "query_task_id": "query-export",
5662            "workflow_type": "counter",
5663            "query_name": "current",
5664            "payload_codec": "json",
5665            "workflow_arguments": {"codec": "json", "blob": "[]"},
5666            "query_arguments": {"codec": "json", "blob": "[]"},
5667            "history_events": [],
5668            "history_export": {
5669                "payloads": {"codec": "json"},
5670                "history_events": [{
5671                    "type": "SignalReceived",
5672                    "payload": {"signal_id": "signal-export", "signal_name": "increment"}
5673                }],
5674                "signals": [{
5675                    "id": "signal-export",
5676                    "name": "increment",
5677                    "status": "applied",
5678                    "workflow_sequence": 1,
5679                    "payload_codec": "json",
5680                    "arguments": "[9]"
5681                }]
5682            }
5683        }))
5684        .expect("query task");
5685
5686        let result = worker.execute_query_task(task).await.expect("query result");
5687        assert_eq!(result, json!(9));
5688    }
5689
5690    #[tokio::test]
5691    async fn query_task_failures_have_stable_reasons() {
5692        let client = Client::new("http://127.0.0.1:8080").expect("client");
5693        let mut worker = Worker::new(client, "rust-workers");
5694        worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5695        worker.register_query(
5696            "counter",
5697            "current",
5698            |_ctx, _args| async move { Ok(json!(0)) },
5699        );
5700
5701        let base_task = QueryTask {
5702            query_task_id: "query-errors".to_string(),
5703            query_task_attempt: 1,
5704            lease_owner: None,
5705            workflow_id: Some("counter-errors".to_string()),
5706            run_id: Some("run-errors".to_string()),
5707            workflow_type: "counter".to_string(),
5708            query_name: "missing".to_string(),
5709            payload_codec: JSON_CODEC.to_string(),
5710            workflow_arguments: Some(json!({"codec": "json", "blob": "[]"})),
5711            query_arguments: Some(json!({"codec": "json", "blob": "[]"})),
5712            history_events: Vec::new(),
5713            history_export: None,
5714            run_status: Some("running".to_string()),
5715        };
5716
5717        let unknown = worker
5718            .execute_query_task(base_task.clone())
5719            .await
5720            .expect_err("unknown query");
5721        assert_eq!(unknown.reason, "rejected_unknown_query");
5722
5723        let mut malformed = base_task;
5724        malformed.query_name = "current".to_string();
5725        malformed.query_arguments = Some(json!({"codec": "json", "blob": "{"}));
5726        let malformed = worker
5727            .execute_query_task(malformed)
5728            .await
5729            .expect_err("malformed payload");
5730        assert_eq!(malformed.reason, "query_payload_decode_failed");
5731
5732        let client = Client::new("http://127.0.0.1:8080").expect("client");
5733        let mut unavailable_worker = Worker::new(client, "rust-workers");
5734        unavailable_worker
5735            .register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5736        let unavailable_task: QueryTask = serde_json::from_value(json!({
5737            "query_task_id": "query-unavailable",
5738            "workflow_type": "counter",
5739            "query_name": "current",
5740            "payload_codec": "json",
5741            "workflow_arguments": {"codec": "json", "blob": "[]"},
5742            "query_arguments": {"codec": "json", "blob": "[]"}
5743        }))
5744        .expect("query task");
5745        let unavailable = unavailable_worker
5746            .execute_query_task(unavailable_task)
5747            .await
5748            .expect_err("query handler unavailable");
5749        assert_eq!(unavailable.reason, "query_handler_unavailable");
5750    }
5751
5752    #[tokio::test]
5753    async fn client_query_decodes_result_and_typed_failure() {
5754        let server = MockWorkerServer::start();
5755        let client = Client::builder(server.base_url())
5756            .timeout(Duration::from_secs(2))
5757            .build()
5758            .expect("client");
5759
5760        let result = client
5761            .query_workflow("counter-1", "current", json!([]))
5762            .await
5763            .expect("query result");
5764        assert_eq!(result, json!({"count": 8}));
5765
5766        let error = client
5767            .query_workflow("counter-1", "missing", json!([]))
5768            .await
5769            .expect_err("unknown query");
5770        let Error::QueryFailed(failure) = error else {
5771            panic!("expected typed query failure");
5772        };
5773        assert_eq!(failure.status, 404);
5774        assert_eq!(failure.reason, "rejected_unknown_query");
5775    }
5776
5777    #[tokio::test]
5778    async fn baseline_worker_endpoints_send_the_baseline_protocol() {
5779        let server = MockWorkerServer::start();
5780        let client = Client::builder(server.base_url())
5781            .timeout(Duration::from_secs(2))
5782            .build()
5783            .expect("client");
5784
5785        client
5786            .register_worker("capture-worker", "capture", vec![], vec![], 1, 1)
5787            .await
5788            .expect("register");
5789        client
5790            .heartbeat_worker("capture-worker", 1, 1)
5791            .await
5792            .expect("heartbeat");
5793        client
5794            .poll_workflow_task("capture-worker", "capture", Duration::from_millis(10))
5795            .await
5796            .expect("workflow poll");
5797        client
5798            .poll_activity_task("capture-worker", "capture", Duration::from_millis(10))
5799            .await
5800            .expect("activity poll");
5801
5802        for path in [
5803            "/api/worker/register",
5804            "/api/worker/heartbeat",
5805            "/api/worker/workflow-tasks/poll",
5806            "/api/worker/activity-tasks/poll",
5807        ] {
5808            assert_eq!(
5809                server.worker_protocol_for(path).as_deref(),
5810                Some(WORKER_PROTOCOL_VERSION),
5811                "unexpected protocol for {path}"
5812            );
5813        }
5814
5815        assert_eq!(
5816            server.request_body("/api/worker/workflow-tasks/poll")["timeout_seconds"],
5817            1
5818        );
5819        assert_eq!(
5820            server.request_body("/api/worker/activity-tasks/poll")["timeout_seconds"],
5821            1
5822        );
5823    }
5824
5825    #[tokio::test]
5826    async fn query_task_endpoints_send_the_query_feature_protocol() {
5827        let server = MockWorkerServer::start();
5828        let client = Client::builder(server.base_url())
5829            .timeout(Duration::from_secs(2))
5830            .build()
5831            .expect("client");
5832
5833        client
5834            .poll_query_task("capture-worker", "capture", Duration::from_millis(10))
5835            .await
5836            .expect("query poll");
5837        client
5838            .complete_query_task("query-capture", "capture-worker", 1, json!(8), JSON_CODEC)
5839            .await
5840            .expect("query complete");
5841        client
5842            .fail_query_task(
5843                "query-capture",
5844                "capture-worker",
5845                1,
5846                "failed",
5847                "query_rejected",
5848                "QueryFailed",
5849            )
5850            .await
5851            .expect("query fail");
5852
5853        for path in [
5854            "/api/worker/query-tasks/poll",
5855            "/api/worker/query-tasks/query-capture/complete",
5856            "/api/worker/query-tasks/query-capture/fail",
5857        ] {
5858            assert_eq!(
5859                server.worker_protocol_for(path).as_deref(),
5860                Some(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION),
5861                "unexpected protocol for {path}"
5862            );
5863        }
5864
5865        assert_eq!(
5866            server.request_body("/api/worker/query-tasks/poll")["timeout_seconds"],
5867            1
5868        );
5869    }
5870
5871    #[tokio::test]
5872    async fn query_protocol_rejection_from_older_server_is_typed() {
5873        let server = MockWorkerServer::reject_query_protocol();
5874        let client = Client::builder(server.base_url())
5875            .timeout(Duration::from_secs(2))
5876            .build()
5877            .expect("client");
5878
5879        let error = client
5880            .poll_query_task("capture-worker", "capture", Duration::from_millis(10))
5881            .await
5882            .expect_err("server below query protocol floor must reject");
5883        let Error::Protocol(failure) = error else {
5884            panic!("expected typed protocol failure");
5885        };
5886
5887        assert_eq!(failure.status, 400);
5888        assert_eq!(failure.reason, "unsupported_protocol_version");
5889        assert_eq!(failure.supported_version.as_deref(), Some("1.7"));
5890        assert_eq!(
5891            failure.requested_version.as_deref(),
5892            Some(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION)
5893        );
5894        assert_eq!(
5895            server
5896                .worker_protocol_for("/api/worker/query-tasks/poll")
5897                .as_deref(),
5898            Some(QUERY_TASK_MINIMUM_WORKER_PROTOCOL_VERSION)
5899        );
5900    }
5901
5902    #[tokio::test]
5903    async fn run_once_without_query_handlers_keeps_pre_query_server_compatibility() {
5904        let server = MockWorkerServer::reject_query_protocol();
5905        let client = Client::builder(server.base_url())
5906            .timeout(Duration::from_secs(2))
5907            .build()
5908            .expect("client");
5909        let mut worker = Worker::new(client, "rust-workers")
5910            .worker_id("baseline-worker")
5911            .poll_timeout(Duration::from_millis(10));
5912
5913        worker.register_workflow("baseline.workflow", |_ctx, _input| async move {
5914            Ok(Value::Null)
5915        });
5916
5917        assert_eq!(worker.run_once().await.expect("baseline run once"), 0);
5918        assert_eq!(
5919            server
5920                .worker_protocol_for("/api/worker/workflow-tasks/poll")
5921                .as_deref(),
5922            Some(WORKER_PROTOCOL_VERSION)
5923        );
5924        assert_eq!(
5925            server.worker_protocol_for("/api/worker/query-tasks/poll"),
5926            None,
5927            "a worker without query handlers must not use the query-task endpoint"
5928        );
5929    }
5930
5931    #[tokio::test]
5932    async fn completion_time_query_rejection_is_typed_without_stopping_worker() {
5933        let server = MockWorkerServer::reject_query_completion();
5934        let client = Client::builder(server.base_url())
5935            .timeout(Duration::from_secs(2))
5936            .build()
5937            .expect("client");
5938
5939        let error = client
5940            .complete_query_task("query-late", "late-worker", 1, json!(8), JSON_CODEC)
5941            .await
5942            .expect_err("expired completion must be rejected");
5943        let Error::QueryFailed(failure) = error else {
5944            panic!("expected typed query failure");
5945        };
5946        assert_eq!(failure.status, 409);
5947        assert_eq!(failure.reason, "query_task_timed_out");
5948
5949        let mut worker = Worker::new(client, "rust-workers")
5950            .worker_id("late-worker")
5951            .poll_timeout(Duration::from_millis(10));
5952        worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
5953        worker.register_query(
5954            "counter",
5955            "current",
5956            |_ctx, _args| async move { Ok(json!(8)) },
5957        );
5958
5959        assert_eq!(worker.run_once().await.expect("late task is handled"), 1);
5960        assert_eq!(
5961            worker
5962                .run_once()
5963                .await
5964                .expect("worker continues after late completion"),
5965            0
5966        );
5967        assert_eq!(
5968            server.request_count("/api/worker/query-tasks/query-late/complete"),
5969            2
5970        );
5971        assert_eq!(
5972            server.request_count("/api/worker/query-tasks/query-late/fail"),
5973            0,
5974            "a server completion rejection must not be reported as an encoding failure"
5975        );
5976    }
5977
5978    #[tokio::test]
5979    async fn activity_only_worker_can_shutdown_without_workflow_poller() {
5980        let server = MockWorkerServer::start();
5981        let client = Client::builder(server.base_url())
5982            .timeout(Duration::from_secs(2))
5983            .build()
5984            .expect("client");
5985        let mut worker = Worker::new(client, "rust-workers")
5986            .worker_id("activity-only-worker")
5987            .poll_timeout(Duration::from_millis(10));
5988
5989        worker.register_activity(
5990            "activity.only",
5991            |_ctx, _args| async move { Ok(Value::Null) },
5992        );
5993
5994        worker.run_until(async {}).await.expect("run worker");
5995    }
5996
5997    #[tokio::test]
5998    async fn workflow_only_worker_can_shutdown_without_activity_poller() {
5999        let server = MockWorkerServer::start();
6000        let client = Client::builder(server.base_url())
6001            .timeout(Duration::from_secs(2))
6002            .build()
6003            .expect("client");
6004        let mut worker = Worker::new(client, "rust-workers")
6005            .worker_id("workflow-only-worker")
6006            .poll_timeout(Duration::from_millis(10));
6007
6008        worker.register_workflow(
6009            "workflow.only",
6010            |_ctx, _input| async move { Ok(Value::Null) },
6011        );
6012
6013        worker.run_until(async {}).await.expect("run worker");
6014    }
6015
6016    #[tokio::test]
6017    async fn worker_heartbeat_observer_receives_server_acknowledgements() {
6018        let server = MockWorkerServer::start();
6019        let client = Client::builder(server.base_url())
6020            .timeout(Duration::from_secs(2))
6021            .build()
6022            .expect("client");
6023        let observations = Arc::new(Mutex::new(Vec::new()));
6024        let observed = Arc::clone(&observations);
6025        let mut worker = Worker::new(client, "rust-workers")
6026            .worker_id("observed-heartbeat-worker")
6027            .poll_timeout(Duration::from_millis(10))
6028            .on_worker_heartbeat(move |observation| {
6029                observed
6030                    .lock()
6031                    .expect("heartbeat observations")
6032                    .push(observation.clone());
6033            });
6034
6035        worker.register_workflow("workflow.observed", |_ctx, _input| async move {
6036            Ok(Value::Null)
6037        });
6038        worker
6039            .run_until(tokio::time::sleep(Duration::from_millis(20)))
6040            .await
6041            .expect("run worker");
6042
6043        let observations = observations.lock().expect("heartbeat observations");
6044        let first = observations.first().expect("heartbeat acknowledgement");
6045        assert_eq!(first.worker_id, "observed-heartbeat-worker");
6046        assert_eq!(first.task_queue, "rust-workers");
6047        assert!(first.acknowledged_at_unix_millis > 0);
6048        assert_eq!(first.acknowledgement, json!({}));
6049    }
6050
6051    #[tokio::test]
6052    async fn worker_retries_poll_and_heartbeat_transport_failures_independently() {
6053        let server = MockWorkerServer::transient_worker_failures();
6054        let client = Client::builder(server.base_url())
6055            .timeout(Duration::from_secs(2))
6056            .build()
6057            .expect("client");
6058        let mut worker = Worker::new(client, "rust-workers")
6059            .worker_id("retry-worker")
6060            .poll_timeout(Duration::from_millis(10))
6061            .retry_policy(WorkerRetryPolicy {
6062                max_retries: 2,
6063                initial_backoff: Duration::from_millis(1),
6064                max_backoff: Duration::from_millis(1),
6065            });
6066        worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
6067        worker.register_activity(
6068            "counter.activity",
6069            |_ctx, _input| async move { Ok(Value::Null) },
6070        );
6071        worker.register_query(
6072            "counter",
6073            "current",
6074            |_ctx, _args| async move { Ok(json!(8)) },
6075        );
6076
6077        worker
6078            .run_until(tokio::time::sleep(Duration::from_millis(75)))
6079            .await
6080            .expect("transient failures must not stop the worker");
6081
6082        for path in [
6083            "/api/worker/heartbeat",
6084            "/api/worker/workflow-tasks/poll",
6085            "/api/worker/activity-tasks/poll",
6086            "/api/worker/query-tasks/poll",
6087        ] {
6088            assert!(
6089                server.request_count(path) >= 2,
6090                "{path} must continue after its transient failure"
6091            );
6092        }
6093    }
6094
6095    #[tokio::test]
6096    async fn worker_bounds_transport_retries() {
6097        let server = MockWorkerServer::unavailable_polls();
6098        let client = Client::builder(server.base_url())
6099            .timeout(Duration::from_secs(2))
6100            .build()
6101            .expect("client");
6102        let mut worker = Worker::new(client, "rust-workers")
6103            .worker_id("bounded-retry-worker")
6104            .poll_timeout(Duration::from_millis(10))
6105            .retry_policy(WorkerRetryPolicy {
6106                max_retries: 2,
6107                initial_backoff: Duration::from_millis(1),
6108                max_backoff: Duration::from_millis(1),
6109            });
6110        worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
6111
6112        let error = worker.run().await.expect_err("retry bound must terminate");
6113        assert!(matches!(error, Error::Transport(_)));
6114        assert_eq!(
6115            server.request_count("/api/worker/workflow-tasks/poll"),
6116            3,
6117            "one initial request plus two retries"
6118        );
6119    }
6120
6121    #[tokio::test]
6122    async fn worker_does_not_retry_authentication_failures() {
6123        let server = MockWorkerServer::unauthorized_polls();
6124        let client = Client::builder(server.base_url())
6125            .timeout(Duration::from_secs(2))
6126            .build()
6127            .expect("client");
6128        let mut worker = Worker::new(client, "rust-workers")
6129            .worker_id("unauthorized-worker")
6130            .poll_timeout(Duration::from_millis(10));
6131        worker.register_workflow("counter", |_ctx, _input| async move { Ok(Value::Null) });
6132
6133        let error = worker
6134            .run()
6135            .await
6136            .expect_err("authentication must terminate");
6137        let Error::Http { status, body } = error else {
6138            panic!("expected stable HTTP authentication error");
6139        };
6140        assert_eq!(status, reqwest::StatusCode::UNAUTHORIZED);
6141        assert!(body.contains("authentication_failed"));
6142        assert_eq!(
6143            server.request_count("/api/worker/workflow-tasks/poll"),
6144            1,
6145            "authentication failures must not be retried"
6146        );
6147    }
6148
6149    #[derive(Clone, Debug)]
6150    struct CapturedRequest {
6151        path: String,
6152        worker_protocol: Option<String>,
6153        body: String,
6154    }
6155
6156    struct MockWorkerServer {
6157        addr: SocketAddr,
6158        stop: Arc<AtomicBool>,
6159        requests: Arc<Mutex<Vec<CapturedRequest>>>,
6160        thread: Option<thread::JoinHandle<()>>,
6161    }
6162
6163    #[derive(Clone, Copy, Default)]
6164    struct MockWorkerBehavior {
6165        reject_query_protocol: bool,
6166        reject_query_completion: bool,
6167        poll_failures_per_path: usize,
6168        heartbeat_failures: usize,
6169        unauthorized_polls: bool,
6170    }
6171
6172    impl MockWorkerServer {
6173        fn start() -> Self {
6174            Self::start_with_behavior(MockWorkerBehavior::default())
6175        }
6176
6177        fn reject_query_protocol() -> Self {
6178            Self::start_with_behavior(MockWorkerBehavior {
6179                reject_query_protocol: true,
6180                ..MockWorkerBehavior::default()
6181            })
6182        }
6183
6184        fn reject_query_completion() -> Self {
6185            Self::start_with_behavior(MockWorkerBehavior {
6186                reject_query_completion: true,
6187                ..MockWorkerBehavior::default()
6188            })
6189        }
6190
6191        fn transient_worker_failures() -> Self {
6192            Self::start_with_behavior(MockWorkerBehavior {
6193                poll_failures_per_path: 1,
6194                heartbeat_failures: 1,
6195                ..MockWorkerBehavior::default()
6196            })
6197        }
6198
6199        fn unavailable_polls() -> Self {
6200            Self::start_with_behavior(MockWorkerBehavior {
6201                poll_failures_per_path: usize::MAX,
6202                ..MockWorkerBehavior::default()
6203            })
6204        }
6205
6206        fn unauthorized_polls() -> Self {
6207            Self::start_with_behavior(MockWorkerBehavior {
6208                unauthorized_polls: true,
6209                ..MockWorkerBehavior::default()
6210            })
6211        }
6212
6213        fn start_with_behavior(behavior: MockWorkerBehavior) -> Self {
6214            let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock server");
6215            listener
6216                .set_nonblocking(true)
6217                .expect("configure mock listener");
6218            let addr = listener.local_addr().expect("mock server address");
6219            let stop = Arc::new(AtomicBool::new(false));
6220            let server_stop = Arc::clone(&stop);
6221            let requests = Arc::new(Mutex::new(Vec::new()));
6222            let server_requests = Arc::clone(&requests);
6223            let thread = thread::spawn(move || {
6224                while !server_stop.load(Ordering::SeqCst) {
6225                    match listener.accept() {
6226                        Ok((mut stream, _)) => {
6227                            handle_mock_worker_request(&mut stream, &server_requests, behavior)
6228                        }
6229                        Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
6230                            thread::sleep(Duration::from_millis(5));
6231                        }
6232                        Err(_) => break,
6233                    }
6234                }
6235            });
6236
6237            Self {
6238                addr,
6239                stop,
6240                requests,
6241                thread: Some(thread),
6242            }
6243        }
6244
6245        fn base_url(&self) -> String {
6246            format!("http://{}", self.addr)
6247        }
6248
6249        fn worker_protocol_for(&self, path: &str) -> Option<String> {
6250            self.requests
6251                .lock()
6252                .expect("captured requests")
6253                .iter()
6254                .find(|request| request.path == path)
6255                .and_then(|request| request.worker_protocol.clone())
6256        }
6257
6258        fn request_count(&self, path: &str) -> usize {
6259            self.requests
6260                .lock()
6261                .expect("captured requests")
6262                .iter()
6263                .filter(|request| request.path == path)
6264                .count()
6265        }
6266
6267        fn request_body(&self, path: &str) -> Value {
6268            let requests = self.requests.lock().expect("captured requests");
6269            let body = &requests
6270                .iter()
6271                .find(|request| request.path == path)
6272                .unwrap_or_else(|| panic!("missing request for {path}"))
6273                .body;
6274            serde_json::from_str(body).unwrap_or_else(|error| {
6275                panic!("invalid JSON request body for {path}: {error}: {body:?}")
6276            })
6277        }
6278    }
6279
6280    impl Drop for MockWorkerServer {
6281        fn drop(&mut self) {
6282            self.stop.store(true, Ordering::SeqCst);
6283            let _ = TcpStream::connect(self.addr);
6284
6285            if let Some(thread) = self.thread.take() {
6286                thread.join().expect("join mock server");
6287            }
6288        }
6289    }
6290
6291    fn handle_mock_worker_request(
6292        stream: &mut TcpStream,
6293        requests: &Arc<Mutex<Vec<CapturedRequest>>>,
6294        behavior: MockWorkerBehavior,
6295    ) {
6296        let _ = stream.set_read_timeout(Some(Duration::from_millis(200)));
6297        let mut buffer = [0_u8; 8192];
6298        let mut request = Vec::new();
6299
6300        loop {
6301            match stream.read(&mut buffer) {
6302                Ok(0) => break,
6303                Ok(read) => {
6304                    request.extend_from_slice(&buffer[..read]);
6305                    if mock_request_is_complete(&request) {
6306                        break;
6307                    }
6308                }
6309                Err(error)
6310                    if matches!(
6311                        error.kind(),
6312                        std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
6313                    ) =>
6314                {
6315                    break;
6316                }
6317                Err(_) => return,
6318            }
6319        }
6320
6321        let request = String::from_utf8_lossy(&request);
6322        let body = request
6323            .split_once("\r\n\r\n")
6324            .map(|(_, body)| body)
6325            .unwrap_or_default();
6326        let path = request
6327            .lines()
6328            .next()
6329            .and_then(|line| line.split_whitespace().nth(1))
6330            .unwrap_or_default();
6331        let worker_protocol = request.lines().find_map(|line| {
6332            let (name, value) = line.split_once(':')?;
6333            name.eq_ignore_ascii_case("X-Durable-Workflow-Protocol-Version")
6334                .then(|| value.trim().to_string())
6335        });
6336        let request_number = {
6337            let mut requests = requests.lock().expect("captured requests");
6338            requests.push(CapturedRequest {
6339                path: path.to_string(),
6340                worker_protocol: worker_protocol.clone(),
6341                body: body.to_string(),
6342            });
6343            requests
6344                .iter()
6345                .filter(|request| request.path == path)
6346                .count()
6347        };
6348
6349        let is_poll = matches!(
6350            path,
6351            "/api/worker/workflow-tasks/poll"
6352                | "/api/worker/activity-tasks/poll"
6353                | "/api/worker/query-tasks/poll"
6354        );
6355        if is_poll && request_number <= behavior.poll_failures_per_path {
6356            return;
6357        }
6358        if path == "/api/worker/heartbeat" && request_number <= behavior.heartbeat_failures {
6359            return;
6360        }
6361        if behavior.unauthorized_polls && is_poll {
6362            write_mock_response(
6363                stream,
6364                "401 Unauthorized",
6365                r#"{"reason":"authentication_failed","message":"invalid worker token"}"#,
6366            );
6367            return;
6368        }
6369
6370        if behavior.reject_query_protocol && path.starts_with("/api/worker/query-tasks/") {
6371            let requested_version = worker_protocol.as_deref().unwrap_or("missing");
6372            let body = format!(
6373                r#"{{"reason":"unsupported_protocol_version","message":"Query tasks require worker protocol 1.8 or newer.","supported_version":"1.7","requested_version":"{requested_version}"}}"#
6374            );
6375            write_mock_response(stream, "400 Bad Request", &body);
6376            return;
6377        }
6378
6379        if behavior.reject_query_completion && path == "/api/worker/query-tasks/query-late/complete"
6380        {
6381            write_mock_response(
6382                stream,
6383                "409 Conflict",
6384                r#"{"reason":"query_task_timed_out","message":"query task timed out before completion"}"#,
6385            );
6386            return;
6387        }
6388
6389        let (status, body) = match path {
6390            "/api/worker/register" => (
6391                "200 OK",
6392                r#"{"worker_id":"mock-worker","registered":true,"heartbeat_interval_seconds":3600}"#,
6393            ),
6394            "/api/worker/heartbeat" => ("200 OK", "{}"),
6395            "/api/worker/activity-tasks/poll" | "/api/worker/workflow-tasks/poll" => {
6396                ("200 OK", r#"{"task":null}"#)
6397            }
6398            "/api/worker/query-tasks/poll"
6399                if behavior.reject_query_completion && request_number == 1 =>
6400            {
6401                (
6402                    "200 OK",
6403                    r#"{"task":{"query_task_id":"query-late","query_task_attempt":1,"lease_owner":"late-worker","workflow_id":"counter-late","run_id":"run-late","workflow_type":"counter","query_name":"current","payload_codec":"json","workflow_arguments":{"codec":"json","blob":"[]"},"query_arguments":{"codec":"json","blob":"[]"},"history_events":[],"run_status":"running"}}"#,
6404                )
6405            }
6406            "/api/worker/query-tasks/poll" => ("200 OK", r#"{"task":null}"#),
6407            "/api/worker/query-tasks/query-capture/complete"
6408            | "/api/worker/query-tasks/query-capture/fail" => ("200 OK", "{}"),
6409            "/api/workflows/counter-1/query/current" => (
6410                "200 OK",
6411                r#"{"workflow_id":"counter-1","query_name":"current","result":{"count":8},"result_envelope":{"codec":"json","blob":"{\"count\":8}"}}"#,
6412            ),
6413            "/api/workflows/counter-1/query/missing" => (
6414                "404 Not Found",
6415                r#"{"workflow_id":"counter-1","query_name":"missing","reason":"rejected_unknown_query","message":"unknown query"}"#,
6416            ),
6417            _ => ("404 Not Found", r#"{"message":"not found"}"#),
6418        };
6419        write_mock_response(stream, status, body);
6420    }
6421
6422    fn mock_request_is_complete(request: &[u8]) -> bool {
6423        let Some(header_end) = request
6424            .windows(4)
6425            .position(|window| window == b"\r\n\r\n")
6426            .map(|position| position + 4)
6427        else {
6428            return false;
6429        };
6430        let headers = String::from_utf8_lossy(&request[..header_end]);
6431        let content_length = headers.lines().find_map(|line| {
6432            let (name, value) = line.split_once(':')?;
6433            name.eq_ignore_ascii_case("content-length")
6434                .then(|| value.trim().parse::<usize>().ok())
6435                .flatten()
6436        });
6437
6438        request.len() >= header_end + content_length.unwrap_or(0)
6439    }
6440
6441    fn write_mock_response(stream: &mut TcpStream, status: &str, body: &str) {
6442        let response = format!(
6443            "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
6444            body.len()
6445        );
6446
6447        let _ = stream.write_all(response.as_bytes());
6448        let _ = stream.flush();
6449    }
6450}