runifold_workflow/
remediation.rs

1use std::{future::Future, pin::Pin, sync::Arc};
2
3use runifold_core::{CapabilitySet, ChildEvent, EventId, RunContext, RunEventKind};
4use serde::{Deserialize, Serialize};
5use serde_json::{Value, json};
6use thiserror::Error;
7
8use crate::checkpoint::WorkflowCheckpointCursor;
9use crate::execution::{
10    check_lifecycle, record_domain, save_checkpoint, validate_exact_usage, validate_usage_floor,
11};
12use crate::workflow::WorkflowNode;
13use crate::{
14    StepId, WorkflowCheckpointPhase, WorkflowCheckpointState, WorkflowError, WorkflowResumePolicy,
15    WorkflowStep,
16};
17
18const MAX_REVIEW_REASON_BYTES: usize = 4_096;
19const MAX_REVIEW_FEEDBACK_BYTES: usize = 1_048_576;
20
21/// Bounded repair policy for one reviewable workflow step.
22#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
23pub struct WorkflowRemediationPolicy {
24    max_repairs: u32,
25}
26
27impl WorkflowRemediationPolicy {
28    /// Creates a policy allowing at most `max_repairs` additional generations.
29    pub const fn new(max_repairs: u32) -> Self {
30        Self { max_repairs }
31    }
32
33    /// Returns the number of additional generations allowed after the first.
34    pub const fn max_repairs(self) -> u32 {
35        self.max_repairs
36    }
37}
38
39/// Canonical request presented to an application-owned output reviewer.
40#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
41pub struct WorkflowReviewRequest {
42    /// Stable workflow node under review.
43    pub step: StepId,
44    /// One-based generation attempt.
45    pub attempt: u32,
46    /// Original value supplied to the repairable node.
47    pub original_input: Value,
48    /// Candidate produced by the current generation attempt.
49    pub candidate: Value,
50}
51
52/// Application-owned decision over one generated candidate.
53#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
54#[serde(tag = "kind", rename_all = "snake_case")]
55#[non_exhaustive]
56pub enum WorkflowReviewVerdict {
57    /// Accept the candidate as the node's stable output.
58    Approve,
59    /// Generate another candidate with structured reviewer feedback.
60    Repair {
61        /// Application-defined findings or repair instructions.
62        feedback: Value,
63    },
64    /// Permanently reject the candidate.
65    Reject {
66        /// Bounded operator-facing explanation.
67        reason: String,
68    },
69}
70
71impl WorkflowReviewVerdict {
72    /// Creates an approval decision.
73    pub const fn approve() -> Self {
74        Self::Approve
75    }
76
77    /// Creates a repair decision with application-defined feedback.
78    ///
79    /// # Errors
80    ///
81    /// Rejects feedback whose canonical JSON representation exceeds 1 MiB.
82    pub fn repair(feedback: Value) -> Result<Self, WorkflowReviewError> {
83        validate_feedback(&feedback)?;
84        Ok(Self::Repair { feedback })
85    }
86
87    /// Creates a permanent rejection with a safe explanation.
88    ///
89    /// # Errors
90    ///
91    /// Rejects blank explanations and values above 4 KiB.
92    pub fn reject(reason: impl Into<String>) -> Result<Self, WorkflowReviewError> {
93        let reason = reason.into();
94        validate_reason(&reason)?;
95        Ok(Self::Reject { reason })
96    }
97
98    fn validate(&self) -> Result<(), WorkflowReviewError> {
99        match self {
100            Self::Approve => Ok(()),
101            Self::Repair { feedback } => validate_feedback(feedback),
102            Self::Reject { reason } => validate_reason(reason),
103        }
104    }
105}
106
107/// Failure returned by an application-owned output reviewer.
108#[derive(Clone, Debug, Error, Eq, PartialEq)]
109#[non_exhaustive]
110pub enum WorkflowReviewError {
111    /// Reviewer feedback exceeded the durable checkpoint limit.
112    #[error("workflow review feedback exceeds the 1 MiB durable limit")]
113    FeedbackTooLarge,
114    /// A rejection reason was blank or exceeded its durable limit.
115    #[error("workflow review rejection reason must contain 1..=4096 bytes")]
116    InvalidRejectionReason,
117    /// The reviewer could not evaluate the candidate.
118    #[error("workflow review failed: {0}")]
119    Execution(String),
120}
121
122/// Boxed asynchronous reviewer result.
123#[cfg(not(target_arch = "wasm32"))]
124pub type WorkflowReviewFuture<'a> =
125    Pin<Box<dyn Future<Output = Result<WorkflowReviewVerdict, WorkflowReviewError>> + Send + 'a>>;
126
127/// Boxed asynchronous reviewer result on single-threaded WASM.
128#[cfg(target_arch = "wasm32")]
129pub type WorkflowReviewFuture<'a> =
130    Pin<Box<dyn Future<Output = Result<WorkflowReviewVerdict, WorkflowReviewError>> + 'a>>;
131
132/// Provider-neutral review boundary for a repairable workflow step.
133pub trait WorkflowReviewer: Send + Sync {
134    /// Reviews one candidate without owning remediation or retry policy.
135    fn review<'a>(
136        &'a self,
137        request: WorkflowReviewRequest,
138        run: &'a RunContext,
139    ) -> WorkflowReviewFuture<'a>;
140}
141
142/// Structured input supplied to the second and later generation attempts.
143///
144/// `input` is a model-ready textual prompt, so an [`crate::AgentStep`] can
145/// consume this value directly. Custom steps can instead inspect the complete
146/// original input, previous candidate, feedback, and attempt fields.
147#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
148pub struct WorkflowRepairInput {
149    /// Model-ready prompt containing trusted runtime repair context.
150    pub input: String,
151    /// Original value supplied to the repairable node.
152    pub original_input: Value,
153    /// Candidate rejected by the preceding review.
154    pub previous_candidate: Value,
155    /// Application-defined reviewer findings or repair instructions.
156    pub feedback: Value,
157    /// One-based generation attempt receiving this input.
158    pub attempt: u32,
159}
160
161impl WorkflowRepairInput {
162    fn new(
163        original_input: Value,
164        previous_candidate: Value,
165        feedback: Value,
166        attempt: u32,
167    ) -> Self {
168        let original_prompt = model_prompt(&original_input);
169        let model_candidate = model_visible_candidate(&previous_candidate);
170        let instruction = json!({
171            "attempt": attempt,
172            "previous_candidate": model_candidate,
173            "feedback": feedback,
174            "instruction": "Produce a corrected candidate that addresses the reviewer feedback. Do not claim the previous candidate was accepted.",
175        });
176        let input = format!(
177            "{original_prompt}\n<runifold_workflow_repair trust=\"runtime\">{instruction}</runifold_workflow_repair>"
178        );
179        Self {
180            input,
181            original_input,
182            previous_candidate,
183            feedback,
184            attempt,
185        }
186    }
187}
188
189/// Durable substate of one repairable workflow node.
190#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
191#[serde(tag = "state", rename_all = "snake_case")]
192#[non_exhaustive]
193pub enum WorkflowRemediationCheckpoint {
194    /// The next generation input is durable and has not started.
195    GenerationReady {
196        /// Canonical input for the pending generation.
197        input: Value,
198    },
199    /// A generation may have partially executed.
200    GenerationInFlight {
201        /// Canonical input supplied to the generation attempt.
202        input: Value,
203    },
204    /// A generated candidate is durable and has not entered review.
205    ReviewReady {
206        /// Durable candidate awaiting review.
207        candidate: Value,
208    },
209    /// Review of a durable candidate may have partially executed.
210    ReviewInFlight {
211        /// Durable candidate supplied to the reviewer.
212        candidate: Value,
213    },
214    /// Review accepted a durable candidate before the outer node commit.
215    Approved {
216        /// Accepted output awaiting the outer node commit.
217        output: Value,
218    },
219    /// Review permanently rejected a candidate.
220    Rejected {
221        /// Candidate rejected by the reviewer.
222        candidate: Value,
223        /// Safe reviewer explanation.
224        reason: String,
225    },
226    /// The configured repair limit was exhausted.
227    Exhausted {
228        /// Last generated candidate.
229        candidate: Value,
230        /// Feedback that would have required another generation.
231        feedback: Value,
232    },
233}
234
235pub(crate) struct RepairableNode {
236    pub(crate) generator: Arc<dyn WorkflowStep>,
237    pub(crate) reviewer: Arc<dyn WorkflowReviewer>,
238    pub(crate) reviewer_capabilities: CapabilitySet,
239    pub(crate) policy: WorkflowRemediationPolicy,
240}
241
242pub(crate) fn prepare_remediation_resume(
243    state: &mut WorkflowCheckpointState,
244    run: &RunContext,
245    policy: WorkflowResumePolicy,
246) -> Result<(), WorkflowError> {
247    let WorkflowCheckpointPhase::Remediating {
248        step,
249        attempt,
250        original_input,
251        checkpoint,
252    } = state.phase.clone()
253    else {
254        return Err(WorkflowError::CheckpointIdentityMismatch);
255    };
256    match checkpoint {
257        WorkflowRemediationCheckpoint::GenerationInFlight { input } => {
258            require_retry_authority(&step, policy)?;
259            validate_usage_floor(state.usage, run.budget().usage())?;
260            state.usage = run.budget().usage();
261            state.phase = WorkflowCheckpointPhase::Remediating {
262                step,
263                attempt,
264                original_input,
265                checkpoint: WorkflowRemediationCheckpoint::GenerationReady { input },
266            };
267        }
268        WorkflowRemediationCheckpoint::ReviewInFlight { candidate } => {
269            require_retry_authority(&step, policy)?;
270            validate_usage_floor(state.usage, run.budget().usage())?;
271            state.usage = run.budget().usage();
272            state.phase = WorkflowCheckpointPhase::Remediating {
273                step,
274                attempt,
275                original_input,
276                checkpoint: WorkflowRemediationCheckpoint::ReviewReady { candidate },
277            };
278        }
279        WorkflowRemediationCheckpoint::Rejected { reason, .. } => {
280            validate_exact_usage(state.usage, run.budget().usage())?;
281            return Err(WorkflowError::RemediationRejected {
282                step,
283                attempts: attempt,
284                reason,
285            });
286        }
287        WorkflowRemediationCheckpoint::Exhausted { .. } => {
288            validate_exact_usage(state.usage, run.budget().usage())?;
289            return Err(WorkflowError::RemediationExhausted {
290                step,
291                attempts: attempt,
292            });
293        }
294        WorkflowRemediationCheckpoint::GenerationReady { .. }
295        | WorkflowRemediationCheckpoint::ReviewReady { .. }
296        | WorkflowRemediationCheckpoint::Approved { .. } => {
297            validate_exact_usage(state.usage, run.budget().usage())?;
298        }
299    }
300    Ok(())
301}
302
303fn require_retry_authority(
304    step: &StepId,
305    policy: WorkflowResumePolicy,
306) -> Result<(), WorkflowError> {
307    if policy == WorkflowResumePolicy::RejectAmbiguous {
308        Err(WorkflowError::AmbiguousCheckpoint { step: step.clone() })
309    } else {
310        Ok(())
311    }
312}
313
314pub(crate) async fn execute_repairable_node(
315    workflow: &str,
316    node: &WorkflowNode,
317    repairable: &RepairableNode,
318    state: &mut WorkflowCheckpointState,
319    run: &RunContext,
320    caused_by: Option<EventId>,
321    checkpoint: &mut Option<&mut WorkflowCheckpointCursor>,
322) -> Result<Value, WorkflowError> {
323    let (mut attempt, original_input, mut remediation) = remediation_state(node, state)?;
324    loop {
325        check_lifecycle(run)?;
326        let context = RemediationAttemptContext {
327            workflow,
328            node,
329            repairable,
330            run,
331            caused_by,
332            attempt,
333            original_input: &original_input,
334        };
335        match remediation {
336            WorkflowRemediationCheckpoint::GenerationReady { input } => {
337                remediation = generate_candidate(&context, state, checkpoint, input).await?;
338            }
339            WorkflowRemediationCheckpoint::ReviewReady { candidate } => {
340                match review_candidate(&context, state, checkpoint, candidate).await? {
341                    RemediationAction::Continue {
342                        next_attempt,
343                        checkpoint,
344                    } => {
345                        attempt = next_attempt;
346                        remediation = checkpoint;
347                    }
348                    RemediationAction::Complete(output) => return Ok(output),
349                }
350            }
351            WorkflowRemediationCheckpoint::Approved { output } => return Ok(output),
352            WorkflowRemediationCheckpoint::Rejected { reason, .. } => {
353                return Err(WorkflowError::RemediationRejected {
354                    step: node.id.clone(),
355                    attempts: attempt,
356                    reason,
357                });
358            }
359            WorkflowRemediationCheckpoint::Exhausted { .. } => {
360                return Err(WorkflowError::RemediationExhausted {
361                    step: node.id.clone(),
362                    attempts: attempt,
363                });
364            }
365            WorkflowRemediationCheckpoint::GenerationInFlight { .. }
366            | WorkflowRemediationCheckpoint::ReviewInFlight { .. } => {
367                return Err(WorkflowError::CheckpointIdentityMismatch);
368            }
369        }
370    }
371}
372
373enum RemediationAction {
374    Continue {
375        next_attempt: u32,
376        checkpoint: WorkflowRemediationCheckpoint,
377    },
378    Complete(Value),
379}
380
381struct RemediationAttemptContext<'a> {
382    workflow: &'a str,
383    node: &'a WorkflowNode,
384    repairable: &'a RepairableNode,
385    run: &'a RunContext,
386    caused_by: Option<EventId>,
387    attempt: u32,
388    original_input: &'a Value,
389}
390
391async fn generate_candidate(
392    context: &RemediationAttemptContext<'_>,
393    state: &mut WorkflowCheckpointState,
394    checkpoint: &mut Option<&mut WorkflowCheckpointCursor>,
395    input: Value,
396) -> Result<WorkflowRemediationCheckpoint, WorkflowError> {
397    state.phase = remediation_phase(
398        context.node,
399        context.attempt,
400        context.original_input.clone(),
401        WorkflowRemediationCheckpoint::GenerationInFlight {
402            input: input.clone(),
403        },
404    );
405    state.usage = context.run.budget().usage();
406    save_checkpoint(checkpoint, state).await?;
407    let candidate = execute_generation(
408        context.workflow,
409        context.node,
410        context.repairable,
411        input,
412        context.attempt,
413        context.run,
414        context.caused_by,
415    )
416    .await?;
417    let remediation = WorkflowRemediationCheckpoint::ReviewReady {
418        candidate: candidate.clone(),
419    };
420    state.phase = remediation_phase(
421        context.node,
422        context.attempt,
423        context.original_input.clone(),
424        remediation.clone(),
425    );
426    state.usage = context.run.budget().usage();
427    save_checkpoint(checkpoint, state).await?;
428    Ok(remediation)
429}
430
431async fn review_candidate(
432    context: &RemediationAttemptContext<'_>,
433    state: &mut WorkflowCheckpointState,
434    checkpoint: &mut Option<&mut WorkflowCheckpointCursor>,
435    candidate: Value,
436) -> Result<RemediationAction, WorkflowError> {
437    state.phase = remediation_phase(
438        context.node,
439        context.attempt,
440        context.original_input.clone(),
441        WorkflowRemediationCheckpoint::ReviewInFlight {
442            candidate: candidate.clone(),
443        },
444    );
445    state.usage = context.run.budget().usage();
446    save_checkpoint(checkpoint, state).await?;
447    let verdict = execute_review(
448        context.workflow,
449        context.node,
450        context.repairable,
451        WorkflowReviewRequest {
452            step: context.node.id.clone(),
453            attempt: context.attempt,
454            original_input: context.original_input.clone(),
455            candidate: candidate.clone(),
456        },
457        context.run,
458        context.caused_by,
459    )
460    .await?;
461    apply_review_verdict(context, state, checkpoint, candidate, verdict).await
462}
463
464async fn apply_review_verdict(
465    context: &RemediationAttemptContext<'_>,
466    state: &mut WorkflowCheckpointState,
467    checkpoint: &mut Option<&mut WorkflowCheckpointCursor>,
468    candidate: Value,
469    verdict: WorkflowReviewVerdict,
470) -> Result<RemediationAction, WorkflowError> {
471    match verdict {
472        WorkflowReviewVerdict::Approve => {
473            state.phase = remediation_phase(
474                context.node,
475                context.attempt,
476                context.original_input.clone(),
477                WorkflowRemediationCheckpoint::Approved {
478                    output: candidate.clone(),
479                },
480            );
481            state.usage = context.run.budget().usage();
482            save_checkpoint(checkpoint, state).await?;
483            Ok(RemediationAction::Complete(candidate))
484        }
485        WorkflowReviewVerdict::Reject { reason } => {
486            state.phase = remediation_phase(
487                context.node,
488                context.attempt,
489                context.original_input.clone(),
490                WorkflowRemediationCheckpoint::Rejected {
491                    candidate,
492                    reason: reason.clone(),
493                },
494            );
495            state.usage = context.run.budget().usage();
496            save_checkpoint(checkpoint, state).await?;
497            Err(WorkflowError::RemediationRejected {
498                step: context.node.id.clone(),
499                attempts: context.attempt,
500                reason,
501            })
502        }
503        WorkflowReviewVerdict::Repair { feedback } => {
504            if context.attempt >= context.repairable.policy.max_repairs().saturating_add(1) {
505                state.phase = remediation_phase(
506                    context.node,
507                    context.attempt,
508                    context.original_input.clone(),
509                    WorkflowRemediationCheckpoint::Exhausted {
510                        candidate,
511                        feedback,
512                    },
513                );
514                state.usage = context.run.budget().usage();
515                save_checkpoint(checkpoint, state).await?;
516                return Err(WorkflowError::RemediationExhausted {
517                    step: context.node.id.clone(),
518                    attempts: context.attempt,
519                });
520            }
521            let next_attempt = context
522                .attempt
523                .checked_add(1)
524                .ok_or(WorkflowError::CheckpointIdentityMismatch)?;
525            let repair_input = WorkflowRepairInput::new(
526                context.original_input.clone(),
527                candidate,
528                feedback,
529                next_attempt,
530            );
531            let remediation = WorkflowRemediationCheckpoint::GenerationReady {
532                input: serde_json::to_value(repair_input)?,
533            };
534            state.phase = remediation_phase(
535                context.node,
536                next_attempt,
537                context.original_input.clone(),
538                remediation.clone(),
539            );
540            state.usage = context.run.budget().usage();
541            save_checkpoint(checkpoint, state).await?;
542            Ok(RemediationAction::Continue {
543                next_attempt,
544                checkpoint: remediation,
545            })
546        }
547    }
548}
549
550fn remediation_state(
551    node: &WorkflowNode,
552    state: &WorkflowCheckpointState,
553) -> Result<(u32, Value, WorkflowRemediationCheckpoint), WorkflowError> {
554    match &state.phase {
555        WorkflowCheckpointPhase::Ready => Ok((
556            1,
557            state.value.clone(),
558            WorkflowRemediationCheckpoint::GenerationReady {
559                input: state.value.clone(),
560            },
561        )),
562        WorkflowCheckpointPhase::Remediating {
563            step,
564            attempt,
565            original_input,
566            checkpoint,
567        } if *step == node.id => Ok((*attempt, original_input.clone(), checkpoint.clone())),
568        _ => Err(WorkflowError::CheckpointIdentityMismatch),
569    }
570}
571
572fn remediation_phase(
573    node: &WorkflowNode,
574    attempt: u32,
575    original_input: Value,
576    checkpoint: WorkflowRemediationCheckpoint,
577) -> WorkflowCheckpointPhase {
578    WorkflowCheckpointPhase::Remediating {
579        step: node.id.clone(),
580        attempt,
581        original_input,
582        checkpoint,
583    }
584}
585
586async fn execute_generation(
587    workflow: &str,
588    node: &WorkflowNode,
589    repairable: &RepairableNode,
590    input: Value,
591    attempt: u32,
592    run: &RunContext,
593    caused_by: Option<EventId>,
594) -> Result<Value, WorkflowError> {
595    let started = record_domain(
596        run,
597        "remediation.generation.started",
598        json!({
599            "workflow": workflow,
600            "step": node.id,
601            "attempt": attempt,
602        }),
603        caused_by,
604    )?;
605    let mut child = run.child(node.capabilities.clone()).map_err(|error| {
606        WorkflowError::AuthorityEscalation {
607            step: node.id.clone(),
608            capability: error.capability,
609        }
610    })?;
611    if let Some(event_id) = started {
612        child = child.with_cause(event_id);
613    }
614    run.record(
615        RunEventKind::Child(ChildEvent::Started {
616            child_run_id: child.run_id(),
617        }),
618        started,
619    )?;
620    match repairable.generator.execute(input, &child).await {
621        Ok(output) => {
622            run.record(
623                RunEventKind::Child(ChildEvent::Completed {
624                    child_run_id: child.run_id(),
625                }),
626                started,
627            )?;
628            record_domain(
629                run,
630                "remediation.generation.completed",
631                json!({
632                    "workflow": workflow,
633                    "step": node.id,
634                    "attempt": attempt,
635                }),
636                started,
637            )?;
638            Ok(output)
639        }
640        Err(source) => {
641            run.record(
642                RunEventKind::Child(ChildEvent::Failed {
643                    child_run_id: child.run_id(),
644                }),
645                started,
646            )?;
647            Err(WorkflowError::Step {
648                step: node.id.clone(),
649                source: Box::new(source),
650            })
651        }
652    }
653}
654
655async fn execute_review(
656    workflow: &str,
657    node: &WorkflowNode,
658    repairable: &RepairableNode,
659    request: WorkflowReviewRequest,
660    run: &RunContext,
661    caused_by: Option<EventId>,
662) -> Result<WorkflowReviewVerdict, WorkflowError> {
663    let attempt = request.attempt;
664    let started = record_domain(
665        run,
666        "remediation.review.started",
667        json!({
668            "workflow": workflow,
669            "step": node.id,
670            "attempt": attempt,
671        }),
672        caused_by,
673    )?;
674    let mut child = run
675        .child(repairable.reviewer_capabilities.clone())
676        .map_err(|error| WorkflowError::AuthorityEscalation {
677            step: node.id.clone(),
678            capability: error.capability,
679        })?;
680    if let Some(event_id) = started {
681        child = child.with_cause(event_id);
682    }
683    run.record(
684        RunEventKind::Child(ChildEvent::Started {
685            child_run_id: child.run_id(),
686        }),
687        started,
688    )?;
689    let verdict = repairable
690        .reviewer
691        .review(request, &child)
692        .await
693        .and_then(|verdict| {
694            verdict.validate()?;
695            Ok(verdict)
696        });
697    match verdict {
698        Ok(verdict) => {
699            run.record(
700                RunEventKind::Child(ChildEvent::Completed {
701                    child_run_id: child.run_id(),
702                }),
703                started,
704            )?;
705            record_domain(
706                run,
707                "remediation.review.completed",
708                json!({
709                    "workflow": workflow,
710                    "step": node.id,
711                    "attempt": attempt,
712                    "verdict": match &verdict {
713                        WorkflowReviewVerdict::Approve => "approve",
714                        WorkflowReviewVerdict::Repair { .. } => "repair",
715                        WorkflowReviewVerdict::Reject { .. } => "reject",
716                    },
717                }),
718                started,
719            )?;
720            Ok(verdict)
721        }
722        Err(source) => {
723            run.record(
724                RunEventKind::Child(ChildEvent::Failed {
725                    child_run_id: child.run_id(),
726                }),
727                started,
728            )?;
729            Err(WorkflowError::Review {
730                step: node.id.clone(),
731                source,
732            })
733        }
734    }
735}
736
737fn model_prompt(input: &Value) -> String {
738    match input {
739        Value::String(prompt) => prompt.clone(),
740        Value::Object(object) => object
741            .get("input")
742            .and_then(Value::as_str)
743            .map_or_else(|| input.to_string(), ToOwned::to_owned),
744        _ => input.to_string(),
745    }
746}
747
748fn model_visible_candidate(candidate: &Value) -> Value {
749    candidate
750        .as_object()
751        .and_then(|object| object.get("input"))
752        .and_then(Value::as_str)
753        .map_or_else(|| candidate.clone(), |text| Value::String(text.to_owned()))
754}
755
756fn validate_feedback(feedback: &Value) -> Result<(), WorkflowReviewError> {
757    if serde_json::to_vec(feedback).is_ok_and(|encoded| encoded.len() <= MAX_REVIEW_FEEDBACK_BYTES)
758    {
759        Ok(())
760    } else {
761        Err(WorkflowReviewError::FeedbackTooLarge)
762    }
763}
764
765fn validate_reason(reason: &str) -> Result<(), WorkflowReviewError> {
766    if reason.trim().is_empty() || reason.len() > MAX_REVIEW_REASON_BYTES {
767        Err(WorkflowReviewError::InvalidRejectionReason)
768    } else {
769        Ok(())
770    }
771}
772
773#[cfg(test)]
774mod tests {
775    use std::sync::{
776        Arc, Mutex,
777        atomic::{AtomicBool, AtomicUsize, Ordering},
778    };
779
780    use futures_executor::block_on;
781    use runifold_core::{
782        Budget, BudgetTracker, CapabilitySet, Checkpoint, CheckpointError, CheckpointErrorKind,
783        CheckpointId, CheckpointStore, InMemoryCheckpointStore, RunContext,
784    };
785    use serde_json::{Value, json};
786
787    use super::{
788        WorkflowRemediationCheckpoint, WorkflowRemediationPolicy, WorkflowRepairInput,
789        WorkflowReviewError, WorkflowReviewFuture, WorkflowReviewRequest, WorkflowReviewVerdict,
790        WorkflowReviewer,
791    };
792    use crate::{
793        Workflow, WorkflowCheckpoint, WorkflowCheckpointPhase, WorkflowCheckpointRevision,
794        WorkflowError, WorkflowForkCommand, WorkflowForkPolicy, WorkflowResumePolicy, WorkflowStep,
795        WorkflowStepFuture,
796    };
797
798    struct CountingGenerator {
799        calls: Arc<AtomicUsize>,
800        inputs: Arc<Mutex<Vec<Value>>>,
801    }
802
803    impl WorkflowStep for CountingGenerator {
804        fn execute<'a>(&'a self, input: Value, _run: &'a RunContext) -> WorkflowStepFuture<'a> {
805            let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
806            self.inputs
807                .lock()
808                .unwrap_or_else(std::sync::PoisonError::into_inner)
809                .push(input);
810            Box::pin(async move { Ok(Value::String(format!("candidate-{attempt}"))) })
811        }
812    }
813
814    struct RepairOnceReviewer {
815        calls: Arc<AtomicUsize>,
816        requests: Arc<Mutex<Vec<WorkflowReviewRequest>>>,
817    }
818
819    impl WorkflowReviewer for RepairOnceReviewer {
820        fn review<'a>(
821            &'a self,
822            request: WorkflowReviewRequest,
823            _run: &'a RunContext,
824        ) -> WorkflowReviewFuture<'a> {
825            self.calls.fetch_add(1, Ordering::SeqCst);
826            let attempt = request.attempt;
827            self.requests
828                .lock()
829                .unwrap_or_else(std::sync::PoisonError::into_inner)
830                .push(request);
831            Box::pin(async move {
832                if attempt == 1 {
833                    WorkflowReviewVerdict::repair(json!({
834                        "code": "unsafe_claim",
835                        "instruction": "remove the unsupported guarantee",
836                    }))
837                } else {
838                    Ok(WorkflowReviewVerdict::approve())
839                }
840            })
841        }
842    }
843
844    struct ApproveReviewer {
845        calls: Arc<AtomicUsize>,
846    }
847
848    impl WorkflowReviewer for ApproveReviewer {
849        fn review<'a>(
850            &'a self,
851            _request: WorkflowReviewRequest,
852            _run: &'a RunContext,
853        ) -> WorkflowReviewFuture<'a> {
854            self.calls.fetch_add(1, Ordering::SeqCst);
855            Box::pin(async { Ok(WorkflowReviewVerdict::approve()) })
856        }
857    }
858
859    struct AlwaysRepairReviewer {
860        calls: Arc<AtomicUsize>,
861    }
862
863    impl WorkflowReviewer for AlwaysRepairReviewer {
864        fn review<'a>(
865            &'a self,
866            request: WorkflowReviewRequest,
867            _run: &'a RunContext,
868        ) -> WorkflowReviewFuture<'a> {
869            self.calls.fetch_add(1, Ordering::SeqCst);
870            Box::pin(
871                async move { WorkflowReviewVerdict::repair(json!({"attempt": request.attempt})) },
872            )
873        }
874    }
875
876    struct RejectReviewer;
877
878    impl WorkflowReviewer for RejectReviewer {
879        fn review<'a>(
880            &'a self,
881            _request: WorkflowReviewRequest,
882            _run: &'a RunContext,
883        ) -> WorkflowReviewFuture<'a> {
884            Box::pin(async { WorkflowReviewVerdict::reject("policy denied the candidate") })
885        }
886    }
887
888    struct FailRevisionOnceStore {
889        inner: InMemoryCheckpointStore,
890        revision: u64,
891        failed: AtomicBool,
892    }
893
894    impl FailRevisionOnceStore {
895        fn new(revision: u64) -> Self {
896            Self {
897                inner: InMemoryCheckpointStore::new(),
898                revision,
899                failed: AtomicBool::new(false),
900            }
901        }
902    }
903
904    impl CheckpointStore for FailRevisionOnceStore {
905        fn load(&self, id: CheckpointId) -> Result<Checkpoint, CheckpointError> {
906            self.inner.load(id)
907        }
908
909        fn compare_and_swap(
910            &self,
911            checkpoint: &Checkpoint,
912            expected_revision: Option<u64>,
913        ) -> Result<(), CheckpointError> {
914            if checkpoint.revision == self.revision && !self.failed.swap(true, Ordering::SeqCst) {
915                return Err(CheckpointError::new(
916                    CheckpointErrorKind::Storage,
917                    "injected remediation checkpoint interruption",
918                ));
919            }
920            self.inner.compare_and_swap(checkpoint, expected_revision)
921        }
922    }
923
924    #[test]
925    fn repairable_step_approves_the_first_candidate() {
926        let generator_calls = Arc::new(AtomicUsize::new(0));
927        let reviewer_calls = Arc::new(AtomicUsize::new(0));
928        let workflow = workflow(
929            generator_calls.clone(),
930            Arc::new(Mutex::new(Vec::new())),
931            ApproveReviewer {
932                calls: reviewer_calls.clone(),
933            },
934            WorkflowRemediationPolicy::new(2),
935        );
936        let run = root_run();
937
938        let outcome = block_on(workflow.run("write a claim", &run)).unwrap();
939
940        assert_eq!(outcome.output, json!("candidate-1"));
941        assert_eq!(generator_calls.load(Ordering::SeqCst), 1);
942        assert_eq!(reviewer_calls.load(Ordering::SeqCst), 1);
943    }
944
945    #[test]
946    fn repair_feedback_is_injected_and_the_new_candidate_is_reviewed() {
947        let generator_calls = Arc::new(AtomicUsize::new(0));
948        let reviewer_calls = Arc::new(AtomicUsize::new(0));
949        let inputs = Arc::new(Mutex::new(Vec::new()));
950        let requests = Arc::new(Mutex::new(Vec::new()));
951        let workflow = workflow(
952            generator_calls.clone(),
953            inputs.clone(),
954            RepairOnceReviewer {
955                calls: reviewer_calls.clone(),
956                requests: requests.clone(),
957            },
958            WorkflowRemediationPolicy::new(2),
959        );
960        let run = root_run();
961
962        let outcome = block_on(workflow.run("write a claim", &run)).unwrap();
963
964        assert_eq!(outcome.output, json!("candidate-2"));
965        assert_eq!(generator_calls.load(Ordering::SeqCst), 2);
966        assert_eq!(reviewer_calls.load(Ordering::SeqCst), 2);
967        let inputs = inputs
968            .lock()
969            .unwrap_or_else(std::sync::PoisonError::into_inner);
970        assert_eq!(inputs[0], json!("write a claim"));
971        let repair: WorkflowRepairInput = serde_json::from_value(inputs[1].clone()).unwrap();
972        assert_eq!(repair.attempt, 2);
973        assert_eq!(repair.original_input, json!("write a claim"));
974        assert_eq!(repair.previous_candidate, json!("candidate-1"));
975        assert_eq!(repair.feedback["code"], "unsafe_claim");
976        assert!(repair.input.contains("runifold_workflow_repair"));
977        let requests = requests
978            .lock()
979            .unwrap_or_else(std::sync::PoisonError::into_inner);
980        assert_eq!(requests[0].attempt, 1);
981        assert_eq!(requests[1].attempt, 2);
982    }
983
984    #[test]
985    fn exhausted_remediation_is_durable_and_does_not_run_again_on_resume() {
986        let generator_calls = Arc::new(AtomicUsize::new(0));
987        let reviewer_calls = Arc::new(AtomicUsize::new(0));
988        let workflow = workflow(
989            generator_calls.clone(),
990            Arc::new(Mutex::new(Vec::new())),
991            AlwaysRepairReviewer {
992                calls: reviewer_calls.clone(),
993            },
994            WorkflowRemediationPolicy::new(1),
995        );
996        let store = Arc::new(InMemoryCheckpointStore::new());
997        let checkpoint = WorkflowCheckpoint::new(store);
998        let run = root_run();
999
1000        let error =
1001            block_on(workflow.run_checkpointed("write a claim", &run, &checkpoint)).unwrap_err();
1002
1003        assert!(matches!(
1004            error,
1005            WorkflowError::RemediationExhausted { attempts: 2, .. }
1006        ));
1007        let (_, state) = checkpoint.load().unwrap();
1008        assert!(matches!(
1009            state.phase,
1010            WorkflowCheckpointPhase::Remediating {
1011                attempt: 2,
1012                checkpoint: WorkflowRemediationCheckpoint::Exhausted { .. },
1013                ..
1014            }
1015        ));
1016        let resumed =
1017            block_on(workflow.resume(&checkpoint, &run, WorkflowResumePolicy::RejectAmbiguous))
1018                .unwrap_err();
1019        assert!(matches!(
1020            resumed,
1021            WorkflowError::RemediationExhausted { attempts: 2, .. }
1022        ));
1023        assert_eq!(generator_calls.load(Ordering::SeqCst), 2);
1024        assert_eq!(reviewer_calls.load(Ordering::SeqCst), 2);
1025    }
1026
1027    #[test]
1028    fn durable_review_ready_resume_does_not_repeat_generation() {
1029        let generator_calls = Arc::new(AtomicUsize::new(0));
1030        let reviewer_calls = Arc::new(AtomicUsize::new(0));
1031        let workflow = workflow(
1032            generator_calls.clone(),
1033            Arc::new(Mutex::new(Vec::new())),
1034            ApproveReviewer {
1035                calls: reviewer_calls.clone(),
1036            },
1037            WorkflowRemediationPolicy::new(1),
1038        );
1039        let store = Arc::new(FailRevisionOnceStore::new(3));
1040        let checkpoint = WorkflowCheckpoint::new(store);
1041        let run = root_run();
1042
1043        let first =
1044            block_on(workflow.run_checkpointed("write a claim", &run, &checkpoint)).unwrap_err();
1045        assert!(matches!(first, WorkflowError::Checkpoint(_)));
1046        assert_eq!(generator_calls.load(Ordering::SeqCst), 1);
1047        assert_eq!(reviewer_calls.load(Ordering::SeqCst), 0);
1048
1049        let outcome =
1050            block_on(workflow.resume(&checkpoint, &run, WorkflowResumePolicy::RejectAmbiguous))
1051                .unwrap();
1052
1053        assert_eq!(outcome.output, json!("candidate-1"));
1054        assert_eq!(generator_calls.load(Ordering::SeqCst), 1);
1055        assert_eq!(reviewer_calls.load(Ordering::SeqCst), 1);
1056    }
1057
1058    #[test]
1059    fn in_flight_generation_requires_explicit_retry_authority() {
1060        let generator_calls = Arc::new(AtomicUsize::new(0));
1061        let reviewer_calls = Arc::new(AtomicUsize::new(0));
1062        let workflow = workflow(
1063            generator_calls.clone(),
1064            Arc::new(Mutex::new(Vec::new())),
1065            ApproveReviewer {
1066                calls: reviewer_calls.clone(),
1067            },
1068            WorkflowRemediationPolicy::new(1),
1069        );
1070        let store = Arc::new(FailRevisionOnceStore::new(2));
1071        let checkpoint = WorkflowCheckpoint::new(store);
1072        let run = root_run();
1073
1074        let first =
1075            block_on(workflow.run_checkpointed("write a claim", &run, &checkpoint)).unwrap_err();
1076        assert!(matches!(first, WorkflowError::Checkpoint(_)));
1077        let rejected =
1078            block_on(workflow.resume(&checkpoint, &run, WorkflowResumePolicy::RejectAmbiguous))
1079                .unwrap_err();
1080        assert!(matches!(
1081            rejected,
1082            WorkflowError::AmbiguousCheckpoint { .. }
1083        ));
1084        let (envelope, _) = checkpoint.load().unwrap();
1085        let rejected_fork = WorkflowForkCommand::new(
1086            checkpoint.id(),
1087            envelope.revision,
1088            WorkflowForkPolicy::RejectAmbiguous,
1089        )
1090        .prepare_checkpoint(envelope.clone())
1091        .unwrap_err();
1092        assert_eq!(rejected_fork.kind, CheckpointErrorKind::Conflict);
1093        let retryable_fork = WorkflowForkCommand::new(
1094            checkpoint.id(),
1095            envelope.revision,
1096            WorkflowForkPolicy::RetryInterruptedStep,
1097        )
1098        .prepare_checkpoint(envelope)
1099        .unwrap();
1100        let revision = WorkflowCheckpointRevision::from_checkpoint(retryable_fork).unwrap();
1101        assert!(matches!(
1102            revision.state.phase,
1103            WorkflowCheckpointPhase::Remediating {
1104                checkpoint: WorkflowRemediationCheckpoint::GenerationReady { .. },
1105                ..
1106            }
1107        ));
1108
1109        let outcome = block_on(workflow.resume(
1110            &checkpoint,
1111            &run,
1112            WorkflowResumePolicy::RetryInterruptedStep,
1113        ))
1114        .unwrap();
1115
1116        assert_eq!(outcome.output, json!("candidate-2"));
1117        assert_eq!(generator_calls.load(Ordering::SeqCst), 2);
1118        assert_eq!(reviewer_calls.load(Ordering::SeqCst), 1);
1119    }
1120
1121    #[test]
1122    fn in_flight_review_requires_explicit_retry_without_repeating_generation() {
1123        let generator_calls = Arc::new(AtomicUsize::new(0));
1124        let reviewer_calls = Arc::new(AtomicUsize::new(0));
1125        let workflow = workflow(
1126            generator_calls.clone(),
1127            Arc::new(Mutex::new(Vec::new())),
1128            ApproveReviewer {
1129                calls: reviewer_calls.clone(),
1130            },
1131            WorkflowRemediationPolicy::new(1),
1132        );
1133        let store = Arc::new(FailRevisionOnceStore::new(4));
1134        let checkpoint = WorkflowCheckpoint::new(store);
1135        let run = root_run();
1136
1137        let first =
1138            block_on(workflow.run_checkpointed("write a claim", &run, &checkpoint)).unwrap_err();
1139        assert!(matches!(first, WorkflowError::Checkpoint(_)));
1140        assert_eq!(generator_calls.load(Ordering::SeqCst), 1);
1141        assert_eq!(reviewer_calls.load(Ordering::SeqCst), 1);
1142        let rejected =
1143            block_on(workflow.resume(&checkpoint, &run, WorkflowResumePolicy::RejectAmbiguous))
1144                .unwrap_err();
1145        assert!(matches!(
1146            rejected,
1147            WorkflowError::AmbiguousCheckpoint { .. }
1148        ));
1149
1150        let outcome = block_on(workflow.resume(
1151            &checkpoint,
1152            &run,
1153            WorkflowResumePolicy::RetryInterruptedStep,
1154        ))
1155        .unwrap();
1156
1157        assert_eq!(outcome.output, json!("candidate-1"));
1158        assert_eq!(generator_calls.load(Ordering::SeqCst), 1);
1159        assert_eq!(reviewer_calls.load(Ordering::SeqCst), 2);
1160    }
1161
1162    #[test]
1163    fn reviewer_rejection_is_a_stable_typed_failure() {
1164        let workflow = workflow(
1165            Arc::new(AtomicUsize::new(0)),
1166            Arc::new(Mutex::new(Vec::new())),
1167            RejectReviewer,
1168            WorkflowRemediationPolicy::new(3),
1169        );
1170
1171        let error = block_on(workflow.run("write a claim", &root_run())).unwrap_err();
1172
1173        assert!(matches!(
1174            error,
1175            WorkflowError::RemediationRejected {
1176                attempts: 1,
1177                reason,
1178                ..
1179            } if reason == "policy denied the candidate"
1180        ));
1181    }
1182
1183    #[test]
1184    fn review_payloads_are_bounded_before_checkpointing() {
1185        assert!(WorkflowReviewVerdict::reject(" ").is_err());
1186        assert!(
1187            WorkflowReviewVerdict::repair(Value::String(
1188                "x".repeat(super::MAX_REVIEW_FEEDBACK_BYTES)
1189            ))
1190            .is_err()
1191        );
1192        assert_eq!(
1193            WorkflowReviewVerdict::reject("x".repeat(super::MAX_REVIEW_REASON_BYTES + 1))
1194                .unwrap_err(),
1195            WorkflowReviewError::InvalidRejectionReason
1196        );
1197    }
1198
1199    #[test]
1200    fn model_repair_prompt_projects_agent_candidate_text_without_full_outcome() {
1201        let repair = WorkflowRepairInput::new(
1202            json!("original prompt"),
1203            json!({
1204                "input": "visible candidate",
1205                "outcome": {"provider_metadata": "not-model-visible"},
1206            }),
1207            json!({"instruction": "fix it"}),
1208            2,
1209        );
1210
1211        assert!(repair.input.contains("visible candidate"));
1212        assert!(!repair.input.contains("not-model-visible"));
1213        assert_eq!(
1214            repair.previous_candidate["outcome"]["provider_metadata"],
1215            "not-model-visible"
1216        );
1217    }
1218
1219    fn workflow<R>(
1220        generator_calls: Arc<AtomicUsize>,
1221        inputs: Arc<Mutex<Vec<Value>>>,
1222        reviewer: R,
1223        policy: WorkflowRemediationPolicy,
1224    ) -> Workflow
1225    where
1226        R: WorkflowReviewer + 'static,
1227    {
1228        Workflow::builder("reviewed-generation")
1229            .repairable_step(
1230                "draft",
1231                CountingGenerator {
1232                    calls: generator_calls,
1233                    inputs,
1234                },
1235                reviewer,
1236                policy,
1237                CapabilitySet::new(),
1238                CapabilitySet::new(),
1239            )
1240            .build()
1241            .unwrap()
1242    }
1243
1244    fn root_run() -> RunContext {
1245        RunContext::root(BudgetTracker::new(Budget::default()), CapabilitySet::new())
1246    }
1247}
runifold_workflow/remediation.rs

runifold_workflow/
remediation.rs
