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