Skip to main content

lash_sansio/sansio/sections/
machine_state.rs

1#[derive(Debug, Serialize, serde::Deserialize)]
2enum MachineState<M: TurnProtocol = UnitTurnProtocol> {
3    PreparingProtocol,
4    WaitingExecutionEnvironment {
5        effect_id: EffectId,
6        update_machine_config: bool,
7    },
8    PrepareIteration,
9    WaitingLlm {
10        effect_id: EffectId,
11        request: Arc<LlmRequest>,
12        driver_state: Option<M::DriverState>,
13    },
14    WaitingTools {
15        effect_id: EffectId,
16        calls: Vec<PendingToolCall>,
17    },
18    WaitingExec {
19        effect_id: EffectId,
20        code: String,
21        driver_state: M::DriverState,
22    },
23    WaitingCheckpoint {
24        effect_id: EffectId,
25        checkpoint: CheckpointKind,
26        on_empty: CheckpointResumeAction,
27    },
28    Finished,
29}
30
31#[derive(Clone, Debug, Serialize, serde::Deserialize)]
32pub struct TurnCheckpoint<M: TurnProtocol = UnitTurnProtocol> {
33    state: MachineState<M>,
34    pending_effects: Vec<Effect<M>>,
35    next_effect_id: u64,
36    #[serde(default)]
37    next_synthetic_message_id: u64,
38    messages: Vec<Message>,
39    events: Vec<SessionEventRecord<M::Event>>,
40    #[serde(default)]
41    turn_causes: Vec<TurnCause>,
42    #[serde(default)]
43    progress_event_cursor: usize,
44    protocol_iteration: usize,
45    protocol_run_offset: usize,
46    cumulative_usage: TokenUsage,
47    termination: TurnTerminationPolicyState,
48    synced_protocol_iteration: Option<usize>,
49}
50
51impl<M: TurnProtocol> Clone for MachineState<M> {
52    fn clone(&self) -> Self {
53        match self {
54            Self::PreparingProtocol => Self::PreparingProtocol,
55            Self::WaitingExecutionEnvironment {
56                effect_id,
57                update_machine_config,
58            } => Self::WaitingExecutionEnvironment {
59                effect_id: *effect_id,
60                update_machine_config: *update_machine_config,
61            },
62            Self::PrepareIteration => Self::PrepareIteration,
63            Self::WaitingLlm {
64                effect_id,
65                request,
66                driver_state,
67            } => Self::WaitingLlm {
68                effect_id: *effect_id,
69                request: Arc::clone(request),
70                driver_state: driver_state.clone(),
71            },
72            Self::WaitingTools { effect_id, calls } => Self::WaitingTools {
73                effect_id: *effect_id,
74                calls: calls.clone(),
75            },
76            Self::WaitingExec {
77                effect_id,
78                code,
79                driver_state,
80            } => Self::WaitingExec {
81                effect_id: *effect_id,
82                code: code.clone(),
83                driver_state: driver_state.clone(),
84            },
85            Self::WaitingCheckpoint {
86                effect_id,
87                checkpoint,
88                on_empty,
89            } => Self::WaitingCheckpoint {
90                effect_id: *effect_id,
91                checkpoint: *checkpoint,
92                on_empty: on_empty.clone(),
93            },
94            Self::Finished => Self::Finished,
95        }
96    }
97}
98
99impl<M: TurnProtocol> MachineState<M> {
100    fn outstanding_effect_id(&self) -> Option<EffectId> {
101        match self {
102            Self::WaitingExecutionEnvironment { effect_id, .. }
103            | Self::WaitingLlm { effect_id, .. }
104            | Self::WaitingTools { effect_id, .. }
105            | Self::WaitingExec { effect_id, .. }
106            | Self::WaitingCheckpoint { effect_id, .. } => Some(*effect_id),
107            Self::PreparingProtocol | Self::PrepareIteration | Self::Finished => None,
108        }
109    }
110
111    fn outstanding_effect(&self) -> Option<Effect<M>> {
112        match self {
113            Self::WaitingExecutionEnvironment {
114                effect_id,
115                update_machine_config,
116            } => Some(Effect::SyncExecutionEnvironment {
117                id: *effect_id,
118                update_machine_config: *update_machine_config,
119            }),
120            Self::WaitingLlm {
121                effect_id, request, ..
122            } => Some(Effect::LlmCall {
123                id: *effect_id,
124                request: Arc::clone(request),
125            }),
126            Self::WaitingTools { effect_id, calls } => Some(Effect::ToolCalls {
127                id: *effect_id,
128                calls: calls.clone(),
129            }),
130            Self::WaitingExec {
131                effect_id, code, ..
132            } => Some(Effect::ExecCode {
133                id: *effect_id,
134                code: code.clone(),
135            }),
136            Self::WaitingCheckpoint {
137                effect_id,
138                checkpoint,
139                ..
140            } => Some(Effect::Checkpoint {
141                id: *effect_id,
142                checkpoint: *checkpoint,
143            }),
144            Self::PreparingProtocol | Self::PrepareIteration | Self::Finished => None,
145        }
146    }
147}
148
149/// Sans-IO state machine for a single session run (multi-turn).
150pub struct TurnMachine<M: TurnProtocol = UnitTurnProtocol> {
151    config: TurnMachineConfig<M>,
152    state: MachineState<M>,
153    pending_effects: VecDeque<Effect<M>>,
154    active_effect_redelivery: bool,
155    next_effect_id: u64,
156    next_synthetic_message_id: u64,
157    messages: MessageSequence,
158    events: Arc<Vec<SessionEventRecord<M::Event>>>,
159    turn_causes: Vec<TurnCause>,
160    progress_event_cursor: usize,
161    protocol_iteration: usize,
162    protocol_run_offset: usize,
163    cumulative_usage: TokenUsage,
164    termination: TurnTerminationPolicyState,
165    synced_protocol_iteration: Option<usize>,
166}