Skip to main content

beam_core/
workflow_actions.rs

1use std::fs;
2use std::path::PathBuf;
3
4use anyhow::Result;
5use serde_json::Value;
6use sha2::{Digest, Sha256};
7
8use crate::{EventDraft, EventLog, WorkflowActor, WorkflowOutputRef};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum WaitKind {
12    HumanGate,
13    Time,
14    Condition,
15}
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum WaitResolution {
19    Approved,
20    Rejected,
21    External,
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum WaitOnTimeout {
26    Fail,
27    Success,
28}
29
30#[derive(Debug, Clone)]
31pub struct CreateWaitInput {
32    pub activity_id: String,
33    pub attempt_id: String,
34    pub node_id: String,
35    pub wait_kind: WaitKind,
36    pub deadline_at: Option<u64>,
37    pub prompt: Option<String>,
38    pub prompt_ref: Option<WorkflowOutputRef>,
39    pub prompt_preview: Option<String>,
40    pub approvers: Option<Vec<String>>,
41    pub on_timeout: Option<WaitOnTimeout>,
42}
43
44#[derive(Debug, Clone)]
45pub struct ResolveWaitInput {
46    pub activity_id: String,
47    pub attempt_id: String,
48    pub resolution: WaitResolution,
49    pub by: String,
50    pub comment: Option<String>,
51    pub output: Option<Value>,
52    pub is_decision_node: bool,
53}
54
55#[derive(Debug, Clone)]
56pub struct ExpireWaitInput {
57    pub activity_id: String,
58    pub attempt_id: String,
59    pub deadline_at: u64,
60    pub exceeded_at_ms: u64,
61    pub on_timeout: Option<WaitOnTimeout>,
62}
63
64#[derive(Debug, Clone)]
65pub struct RequestCancelInput {
66    pub target: Value,
67    pub reason: String,
68    pub by: String,
69}
70
71#[derive(Debug, Clone)]
72pub struct DeliverCancelInput {
73    pub target: Value,
74    pub activity_id: String,
75}
76
77#[derive(Debug, Clone)]
78pub struct CompleteActivityCancelInput {
79    pub activity_id: String,
80    pub attempt_id: String,
81    pub cancel_origin_event_id: String,
82}
83
84#[derive(Debug, Clone)]
85pub struct CompleteNodeCancelInput {
86    pub node_id: String,
87    pub cancel_origin_event_id: String,
88}
89
90#[derive(Debug, Clone)]
91pub struct CompleteRunCancelInput {
92    pub cancel_origin_event_id: String,
93}
94
95pub async fn create_wait(log: &mut EventLog, input: CreateWaitInput) -> Result<Value> {
96    Ok(serde_json::to_value(log.append(EventDraft {
97        event_type: "waitCreated".to_string(),
98        actor: WorkflowActor::Scheduler,
99        payload: serde_json::json!({
100            "activityId": input.activity_id,
101            "attemptId": input.attempt_id,
102            "nodeId": input.node_id,
103            "waitKind": wait_kind_str(input.wait_kind),
104            "deadlineAt": input.deadline_at,
105            "prompt": input.prompt,
106            "promptRef": input.prompt_ref,
107            "promptPreview": input.prompt_preview,
108            "approvers": input.approvers,
109            "onTimeout": input.on_timeout.map(wait_timeout_str),
110        }),
111        timestamp: None,
112        payload_hash: None,
113    })?)?)
114}
115
116pub async fn resolve_wait(log: &mut EventLog, input: ResolveWaitInput) -> Result<Value> {
117    let mut external_refs = serde_json::Map::from_iter([
118        (
119            "resolution".to_string(),
120            Value::String(wait_resolution_str(input.resolution).to_string()),
121        ),
122        ("by".to_string(), Value::String(input.by.clone())),
123        (
124            "comment".to_string(),
125            Value::String(input.comment.clone().unwrap_or_default()),
126        ),
127    ]);
128    if let Some(extra) = input.output.clone() {
129        match extra {
130            Value::Object(map) => {
131                for (key, value) in map {
132                    external_refs.insert(key, value);
133                }
134            }
135            value => {
136                external_refs.insert("output".to_string(), value);
137            }
138        }
139    }
140    let output_ref = write_json_blob(log, Value::Object(external_refs.clone()))?;
141    let resolution_event = log.append(EventDraft {
142        event_type: "waitResolved".to_string(),
143        actor: match input.resolution {
144            WaitResolution::External => WorkflowActor::System,
145            _ => WorkflowActor::Human,
146        },
147        payload: serde_json::json!({
148            "activityId": input.activity_id,
149            "resolution": wait_resolution_str(input.resolution),
150            "by": input.by,
151            "comment": input.comment,
152        }),
153        timestamp: None,
154        payload_hash: None,
155    })?;
156
157    let terminal = match input.resolution {
158        WaitResolution::Rejected if !input.is_decision_node => log.append(EventDraft {
159            event_type: "activityFailed".to_string(),
160            actor: WorkflowActor::Scheduler,
161            payload: serde_json::json!({
162                "activityId": input.activity_id,
163                "attemptId": input.attempt_id,
164                "error": {
165                    "errorCode": "InputValidationFailed",
166                    "errorClass": "userFault",
167                    "errorMessage": format!(
168                        "Wait resolved with rejected by {}{}",
169                        input.by,
170                        input.comment.as_ref().map(|c| format!(": {}", c)).unwrap_or_default()
171                    ),
172                }
173            }),
174            timestamp: None,
175            payload_hash: None,
176        })?,
177        _ => log.append(EventDraft {
178            event_type: "activitySucceeded".to_string(),
179            actor: WorkflowActor::Scheduler,
180            payload: serde_json::json!({
181                "activityId": input.activity_id,
182                "attemptId": input.attempt_id,
183                "outputRef": output_ref,
184                "externalRefs": external_refs,
185            }),
186            timestamp: None,
187            payload_hash: None,
188        })?,
189    };
190
191    Ok(serde_json::json!({
192        "resolutionEventId": resolution_event.event_id,
193        "terminalEventId": terminal.event_id,
194    }))
195}
196
197pub async fn expire_wait(log: &mut EventLog, input: ExpireWaitInput) -> Result<Value> {
198    let _deadline_event = log.append(EventDraft {
199        event_type: "waitDeadlineExceeded".to_string(),
200        actor: WorkflowActor::Scheduler,
201        payload: serde_json::json!({
202            "activityId": input.activity_id,
203            "deadlineAt": input.deadline_at,
204            "exceededAtMs": input.exceeded_at_ms,
205        }),
206        timestamp: None,
207        payload_hash: None,
208    })?;
209    let terminal = match input.on_timeout.unwrap_or(WaitOnTimeout::Fail) {
210        WaitOnTimeout::Fail => log.append(EventDraft {
211            event_type: "activityFailed".to_string(),
212            actor: WorkflowActor::Scheduler,
213            payload: serde_json::json!({
214                "activityId": input.activity_id,
215                "attemptId": input.attempt_id,
216                "error": {
217                    "errorCode": "WaitDeadlineExceeded",
218                    "errorClass": "userFault",
219                    "errorMessage": "Wait deadline exceeded"
220                }
221            }),
222            timestamp: None,
223            payload_hash: None,
224        })?,
225        WaitOnTimeout::Success => log.append(EventDraft {
226            event_type: "activitySucceeded".to_string(),
227            actor: WorkflowActor::Scheduler,
228            payload: serde_json::json!({
229                "activityId": input.activity_id,
230                "attemptId": input.attempt_id,
231                "externalRefs": {
232                    "defaultedToTimeout": true,
233                }
234            }),
235            timestamp: None,
236            payload_hash: None,
237        })?,
238    };
239    Ok(serde_json::json!({
240        "deadlineEventId": "todo",
241        "terminalEventId": terminal.event_id,
242    }))
243}
244
245pub async fn request_cancel(
246    log: &mut EventLog,
247    input: RequestCancelInput,
248    actor: WorkflowActor,
249) -> Result<Value> {
250    let event = log.append(EventDraft {
251        event_type: "cancelRequested".to_string(),
252        actor,
253        payload: serde_json::json!({
254            "target": input.target,
255            "reason": input.reason,
256            "by": input.by,
257        }),
258        timestamp: None,
259        payload_hash: None,
260    })?;
261    Ok(serde_json::json!({ "eventId": event.event_id }))
262}
263
264pub async fn deliver_cancel(
265    log: &mut EventLog,
266    input: DeliverCancelInput,
267    actor: WorkflowActor,
268) -> Result<Value> {
269    let event = log.append(EventDraft {
270        event_type: "cancelDelivered".to_string(),
271        actor,
272        payload: serde_json::json!({
273            "target": input.target,
274            "activityId": input.activity_id,
275        }),
276        timestamp: None,
277        payload_hash: None,
278    })?;
279    Ok(serde_json::json!({ "eventId": event.event_id }))
280}
281
282pub async fn complete_activity_cancel(
283    log: &mut EventLog,
284    input: CompleteActivityCancelInput,
285    actor: WorkflowActor,
286) -> Result<Value> {
287    let event = log.append(EventDraft {
288        event_type: "activityCanceled".to_string(),
289        actor,
290        payload: serde_json::json!({
291            "activityId": input.activity_id,
292            "attemptId": input.attempt_id,
293            "cancelOriginEventId": input.cancel_origin_event_id,
294        }),
295        timestamp: None,
296        payload_hash: None,
297    })?;
298    Ok(serde_json::json!({ "eventId": event.event_id }))
299}
300
301pub async fn complete_node_cancel(
302    log: &mut EventLog,
303    input: CompleteNodeCancelInput,
304    actor: WorkflowActor,
305) -> Result<Value> {
306    let event = log.append(EventDraft {
307        event_type: "nodeCanceled".to_string(),
308        actor,
309        payload: serde_json::json!({
310            "nodeId": input.node_id,
311            "cancelOriginEventId": input.cancel_origin_event_id,
312        }),
313        timestamp: None,
314        payload_hash: None,
315    })?;
316    Ok(serde_json::json!({ "eventId": event.event_id }))
317}
318
319pub async fn complete_run_cancel(
320    log: &mut EventLog,
321    input: CompleteRunCancelInput,
322    actor: WorkflowActor,
323) -> Result<Value> {
324    let event = log.append(EventDraft {
325        event_type: "runCanceled".to_string(),
326        actor,
327        payload: serde_json::json!({
328            "cancelOriginEventId": input.cancel_origin_event_id,
329        }),
330        timestamp: None,
331        payload_hash: None,
332    })?;
333    Ok(serde_json::json!({ "eventId": event.event_id }))
334}
335
336fn wait_kind_str(kind: WaitKind) -> &'static str {
337    match kind {
338        WaitKind::HumanGate => "human-gate",
339        WaitKind::Time => "time",
340        WaitKind::Condition => "condition",
341    }
342}
343
344fn wait_resolution_str(resolution: WaitResolution) -> &'static str {
345    match resolution {
346        WaitResolution::Approved => "approved",
347        WaitResolution::Rejected => "rejected",
348        WaitResolution::External => "external",
349    }
350}
351
352fn wait_timeout_str(timeout: WaitOnTimeout) -> &'static str {
353    match timeout {
354        WaitOnTimeout::Fail => "fail",
355        WaitOnTimeout::Success => "success",
356    }
357}
358
359fn write_json_blob(log: &mut EventLog, value: Value) -> Result<WorkflowOutputRef> {
360    let bytes = serde_json::to_vec(&value)?;
361    let hash = sha256_hex(&bytes);
362    let path = PathBuf::from(&log.blob_dir).join(&hash);
363    fs::write(&path, &bytes)?;
364    Ok(WorkflowOutputRef {
365        output_hash: format!("sha256:{hash}"),
366        output_path: path.display().to_string(),
367        output_bytes: bytes.len(),
368        output_schema_version: 1,
369        content_type: Some("application/json".to_string()),
370    })
371}
372
373fn sha256_hex(bytes: &[u8]) -> String {
374    let mut hasher = Sha256::new();
375    hasher.update(bytes);
376    lower_hex(&hasher.finalize())
377}
378
379fn lower_hex(bytes: &[u8]) -> String {
380    const HEX: &[u8; 16] = b"0123456789abcdef";
381    let mut out = String::with_capacity(bytes.len() * 2);
382    for byte in bytes {
383        out.push(HEX[(byte >> 4) as usize] as char);
384        out.push(HEX[(byte & 0x0f) as usize] as char);
385    }
386    out
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392    use crate::BeamPaths;
393    use std::time::{SystemTime, UNIX_EPOCH};
394
395    fn temp_paths(label: &str) -> BeamPaths {
396        BeamPaths::from_root(std::env::temp_dir().join(format!(
397            "beam-workflow-actions-{label}-{}-{}",
398            std::process::id(),
399            SystemTime::now()
400                .duration_since(UNIX_EPOCH)
401                .unwrap_or_default()
402                .as_nanos()
403        )))
404    }
405
406    #[tokio::test]
407    async fn cancel_and_wait_helpers_write_expected_events() {
408        let paths = temp_paths("write");
409        let mut log = EventLog::new("run-1", paths.workflow_runs_dir()).expect("log");
410        let wait = create_wait(
411            &mut log,
412            CreateWaitInput {
413                activity_id: "run-1::gate::node-a".to_string(),
414                attempt_id: "run-1::gate::node-a::att-1".to_string(),
415                node_id: "node-a".to_string(),
416                wait_kind: WaitKind::HumanGate,
417                deadline_at: None,
418                prompt: Some("approve?".to_string()),
419                prompt_ref: None,
420                prompt_preview: None,
421                approvers: None,
422                on_timeout: Some(WaitOnTimeout::Fail),
423            },
424        )
425        .await
426        .expect("create wait");
427        assert_eq!(wait["event_id"], "run-1-1");
428
429        let cancel = request_cancel(
430            &mut log,
431            RequestCancelInput {
432                target: serde_json::json!({ "kind": "run", "runId": "run-1" }),
433                reason: "stop".to_string(),
434                by: "tester".to_string(),
435            },
436            WorkflowActor::Human,
437        )
438        .await
439        .expect("request cancel");
440        assert_eq!(cancel["eventId"], "run-1-2");
441
442        let events = log.read_all().expect("events");
443        assert_eq!(events.len(), 2);
444        assert_eq!(events[0].event_type, "waitCreated");
445        assert_eq!(events[1].event_type, "cancelRequested");
446
447        let resolved = resolve_wait(
448            &mut log,
449            ResolveWaitInput {
450                activity_id: "run-1::gate::node-a".to_string(),
451                attempt_id: "run-1::gate::node-a::att-1".to_string(),
452                resolution: WaitResolution::Approved,
453                by: "tester".to_string(),
454                comment: Some("ok".to_string()),
455                output: None,
456                is_decision_node: false,
457            },
458        )
459        .await
460        .expect("resolve");
461        assert_eq!(resolved["resolutionEventId"], "run-1-3");
462    }
463
464    #[tokio::test]
465    async fn decision_node_reject_resolves_to_success() {
466        let paths = temp_paths("decision");
467        let mut log = EventLog::new("run-2", paths.workflow_runs_dir()).expect("log");
468        let _ = create_wait(
469            &mut log,
470            CreateWaitInput {
471                activity_id: "run-2::work::node-decision::att-1".to_string(),
472                attempt_id: "run-2::work::node-decision::att-1".to_string(),
473                node_id: "node-decision".to_string(),
474                wait_kind: WaitKind::HumanGate,
475                deadline_at: None,
476                prompt: Some("choose".to_string()),
477                prompt_ref: None,
478                prompt_preview: None,
479                approvers: None,
480                on_timeout: Some(WaitOnTimeout::Fail),
481            },
482        )
483        .await
484        .expect("create wait");
485        let _ = resolve_wait(
486            &mut log,
487            ResolveWaitInput {
488                activity_id: "run-2::work::node-decision::att-1".to_string(),
489                attempt_id: "run-2::work::node-decision::att-1".to_string(),
490                resolution: WaitResolution::Rejected,
491                by: "tester".to_string(),
492                comment: Some("nope".to_string()),
493                output: None,
494                is_decision_node: true,
495            },
496        )
497        .await
498        .expect("resolve");
499        let events = log.read_all().expect("events");
500        assert_eq!(events[1].event_type, "waitResolved");
501        assert_eq!(events[2].event_type, "activitySucceeded");
502        let _ = std::fs::remove_dir_all(paths.root());
503    }
504}