taquba-workflow 0.11.0

Durable, at-least-once workflow runtime on top of the Taquba task queue. Particularly well-suited for AI agent runs.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186

//! The in-process registry of active runs. Every method takes and
//! releases the lock internally, no method awaits and the guard never
//! leaves this module, so the registry lock cannot be held across an
//! await point.

use std::collections::HashMap;
use std::sync::Mutex;

use crate::runtime::{RunState, RunStatus};

/// Per-active-run state retained by the runtime. Combines the publicly
/// observable [`RunStatus`] with the in-process state needed to resolve
/// [`WorkflowRuntime::cancel`](crate::WorkflowRuntime::cancel) races:
/// the Taquba job currently representing the run (so `cancel` can
/// target it), the submitter's headers (so the notification includes
/// the right metadata even when `cancel` terminates a pending step
/// directly) and a flag for any pending cancellation request.
struct RegistryEntry {
    status: RunStatus,
    current_job_id: String,
    user_headers: HashMap<String, String>,
    cancel_requested: bool,
    /// SHA-256 of the original `spec.input`. `Some` for entries created
    /// by [`WorkflowRuntime::submit`](crate::WorkflowRuntime::submit);
    /// `None` for entries created by a worker resuming a step after
    /// restart, which does not have access to the original input. The
    /// duplicate-submit check falls through to the durable record
    /// (which always includes the hash) when this is `None`.
    input_hash: Option<[u8; 32]>,
}

/// The map of active runs, keyed by run id. Terminal runs are removed
/// by [`Self::forget`]; the map holds only runs this process considers
/// live.
#[derive(Default)]
pub(crate) struct RunRegistry {
    runs: Mutex<HashMap<String, RegistryEntry>>,
}

impl RunRegistry {
    /// The stored input hash of `run_id`, when an entry created by
    /// `submit` exists. `None` for an unknown run and for a
    /// worker-resumed entry, which stores no hash.
    pub(crate) fn known_input_hash(&self, run_id: &str) -> Option<[u8; 32]> {
        self.runs
            .lock()
            .unwrap()
            .get(run_id)
            .and_then(|entry| entry.input_hash)
    }

    /// The id of the queue job backing the run's current step, when the
    /// run is tracked.
    pub(crate) fn current_job_id(&self, run_id: &str) -> Option<String> {
        self.runs
            .lock()
            .unwrap()
            .get(run_id)
            .map(|entry| entry.current_job_id.clone())
    }

    /// Record a newly submitted run: step 0, [`RunState::Pending`],
    /// backed by `job_id`.
    pub(crate) fn insert_submitted(
        &self,
        run_id: &str,
        job_id: &str,
        user_headers: HashMap<String, String>,
        input_hash: [u8; 32],
    ) {
        self.runs.lock().unwrap().insert(
            run_id.to_string(),
            RegistryEntry {
                status: RunStatus {
                    run_id: run_id.to_string(),
                    state: RunState::Pending,
                    current_step: 0,
                },
                current_job_id: job_id.to_string(),
                user_headers,
                cancel_requested: false,
                input_hash: Some(input_hash),
            },
        );
    }

    /// The status of `run_id`, with a pending cancellation reported as
    /// [`RunState::Cancelling`] regardless of the underlying step
    /// lifecycle position.
    pub(crate) fn status(&self, run_id: &str) -> Option<RunStatus> {
        self.runs.lock().unwrap().get(run_id).map(|entry| {
            let mut status = entry.status.clone();
            if entry.cancel_requested {
                status.state = RunState::Cancelling;
            }
            status
        })
    }

    /// Whether a cancellation has been requested for `run_id`.
    pub(crate) fn cancel_requested(&self, run_id: &str) -> bool {
        self.runs
            .lock()
            .unwrap()
            .get(run_id)
            .is_some_and(|entry| entry.cancel_requested)
    }

    /// Set the cancellation flag on `run_id` and return the job id,
    /// submitter headers and current step the cancellation targets.
    /// `None` when the run is not active in this registry.
    pub(crate) fn request_cancel(
        &self,
        run_id: &str,
    ) -> Option<(String, HashMap<String, String>, u32)> {
        let mut runs = self.runs.lock().unwrap();
        let entry = runs.get_mut(run_id)?;
        entry.cancel_requested = true;
        Some((
            entry.current_job_id.clone(),
            entry.user_headers.clone(),
            entry.status.current_step,
        ))
    }

    /// Transition the entry for `run_id` into [`RunState::Running`] for
    /// `step_number`, recording the Taquba job id backing the step so a
    /// concurrent cancellation can target it. Creates a fresh entry,
    /// with no input hash and no cancellation flag, when the run is
    /// unknown to this registry (a worker resuming after a restart
    /// first learns of the run by claiming its step).
    pub(crate) fn mark_running(
        &self,
        run_id: &str,
        step_number: u32,
        job_id: &str,
        user_headers: &HashMap<String, String>,
    ) {
        let mut runs = self.runs.lock().unwrap();
        match runs.get_mut(run_id) {
            Some(entry) => {
                entry.status.state = RunState::Running;
                entry.status.current_step = step_number;
                entry.current_job_id = job_id.to_string();
            }
            None => {
                runs.insert(
                    run_id.to_string(),
                    RegistryEntry {
                        status: RunStatus {
                            run_id: run_id.to_string(),
                            state: RunState::Running,
                            current_step: step_number,
                        },
                        current_job_id: job_id.to_string(),
                        user_headers: user_headers.clone(),
                        cancel_requested: false,
                        input_hash: None,
                    },
                );
            }
        }
    }

    /// Transition the entry for `run_id` to [`RunState::Pending`] at
    /// `next_step`, backed by `next_job_id`. The cancellation flag is
    /// deliberately preserved: a cancellation issued during the
    /// just-settled step must survive the advance.
    pub(crate) fn mark_pending(&self, run_id: &str, next_step: u32, next_job_id: String) {
        if let Some(entry) = self.runs.lock().unwrap().get_mut(run_id) {
            entry.status.state = RunState::Pending;
            entry.status.current_step = next_step;
            entry.current_job_id = next_job_id;
        }
    }

    /// Remove the entry for `run_id`.
    pub(crate) fn forget(&self, run_id: &str) {
        self.runs.lock().unwrap().remove(run_id);
    }

    #[cfg(test)]
    pub(crate) fn is_empty(&self) -> bool {
        self.runs.lock().unwrap().is_empty()
    }
}