cookie_agent_protocol/
session_model.rs

1//! Session-local model continuation derived from existing durable events.
2use crate::{EventPayload, ModelFinishReason, ModelSelection, RunId, RunSelection, StoredEvent};
3
4#[derive(Clone, Debug, Default, PartialEq)]
5pub struct SessionModelState {
6    pub selection: Option<RunSelection>,
7    pub successful_fallback: bool,
8    run: Option<RunId>,
9    chain: Vec<ModelSelection>,
10}
11
12impl SessionModelState {
13    pub fn from_events(events: &[StoredEvent]) -> Self {
14        let mut state = Self::default();
15        for event in events {
16            state.apply(event.run_id, &event.payload);
17        }
18        state
19    }
20
21    pub fn apply(&mut self, run: Option<RunId>, event: &EventPayload) {
22        match event {
23            EventPayload::SessionCreated {
24                creation_selection, ..
25            } => {
26                self.selection = Some(creation_selection.clone());
27            }
28            EventPayload::RunStarted {
29                selection,
30                selected_suffix,
31                ..
32            } => {
33                self.successful_fallback &= self.selection.as_ref() == Some(selection);
34                self.selection = Some(selection.clone());
35                self.chain = selected_suffix
36                    .iter()
37                    .map(|binding| binding.selection.clone())
38                    .collect();
39                self.run = run;
40            }
41            EventPayload::ModelTurnCommitted {
42                resolved_model,
43                turn,
44                ..
45            } if run.is_some() && run == self.run && successful_finish(&turn.finish_reason) => {
46                // Temporary skill overrides outside the frozen chain do not change
47                // the session default. ModelFallback alone is only an attempted move.
48                if self
49                    .chain
50                    .iter()
51                    .position(|model| *model == resolved_model.selection)
52                    .is_some_and(|index| index > 0)
53                    && let Some(selection) = &mut self.selection
54                {
55                    selection.model = resolved_model.selection.clone();
56                    self.successful_fallback = true;
57                }
58            }
59            _ => {}
60        }
61    }
62
63    /// Resolve a stale/default selection against a freshly frozen chain. Exact
64    /// model+variant identities, not an index from an earlier configuration, match.
65    pub fn continuation(
66        &self,
67        requested: &RunSelection,
68        chain: &[ModelSelection],
69        reset: bool,
70    ) -> RunSelection {
71        if !reset
72            && self.successful_fallback
73            && let Some(selected) = &self.selection
74            && selected.agent == requested.agent
75            && selected.preset == requested.preset
76            && let Some(requested_index) = chain.iter().position(|model| *model == requested.model)
77            && let Some(selected_index) = chain.iter().position(|model| *model == selected.model)
78            && requested_index < selected_index
79        {
80            return selected.clone();
81        }
82        requested.clone()
83    }
84}
85
86fn successful_finish(reason: &ModelFinishReason) -> bool {
87    match reason {
88        // A tool-calling or output-limited turn still establishes a working model.
89        ModelFinishReason::Stop | ModelFinishReason::ToolCalls | ModelFinishReason::Length => true,
90        ModelFinishReason::ContentFilter | ModelFinishReason::Cancelled | ModelFinishReason::Error
91        | ModelFinishReason::Aborted | ModelFinishReason::Timeout | ModelFinishReason::Refused
92        // Unknown and vendor-specific reasons are not positive success evidence.
93        | ModelFinishReason::Unknown | ModelFinishReason::Other(_) => false,
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    fn selection(model: &str) -> RunSelection {
102        RunSelection {
103            agent: "primary".parse().unwrap(),
104            preset: None,
105            model: ModelSelection {
106                model: format!("test/{model}").parse().unwrap(),
107                variant: None,
108            },
109        }
110    }
111
112    fn committed(model: &RunSelection, finish_reason: ModelFinishReason) -> EventPayload {
113        EventPayload::ModelTurnCommitted {
114            attempt_id: crate::AttemptId::new_v7(),
115            model_turn_seq: 2,
116            input_through_seq: 1,
117            resolved_model: crate::ResolvedModelRef {
118                selection: model.model.clone(),
119                provider_id: "test".parse().unwrap(),
120                model_id: model.model.model.model_id().clone(),
121                adapter_id: crate::AdaptorId::OpenaiCompatible,
122                selection_fingerprint: crate::Sha256Digest::of_bytes(b"fixture"),
123            },
124            turn: crate::PersistedModelTurn {
125                content: vec![],
126                provider_options: Default::default(),
127                finish_reason,
128                usage: Default::default(),
129                response_metadata: Default::default(),
130                provider_metadata: Default::default(),
131                native_replay: None,
132            },
133            warnings: vec![],
134        }
135    }
136
137    #[test]
138    fn only_successful_finish_reasons_advance_fallback_and_later_failure_keeps_progress() {
139        let a = selection("a");
140        let b = selection("b");
141        let c = selection("c");
142        let run = Some(RunId::new_v7());
143        let initial = SessionModelState {
144            selection: Some(a.clone()),
145            run,
146            chain: vec![a.model.clone(), b.model.clone(), c.model.clone()],
147            successful_fallback: false,
148        };
149        for (reason, success) in [
150            (ModelFinishReason::Stop, true),
151            (ModelFinishReason::ToolCalls, true),
152            (ModelFinishReason::Length, true),
153            (ModelFinishReason::Cancelled, false),
154            (ModelFinishReason::Aborted, false),
155            (ModelFinishReason::Error, false),
156            (ModelFinishReason::Timeout, false),
157            (ModelFinishReason::ContentFilter, false),
158            (ModelFinishReason::Refused, false),
159            (ModelFinishReason::Unknown, false),
160            (ModelFinishReason::Other("vendor-specific".into()), false),
161        ] {
162            let mut state = initial.clone();
163            state.apply(run, &committed(&b, reason.clone()));
164            assert_eq!(
165                state.selection,
166                Some(if success { b.clone() } else { a.clone() }),
167                "{reason:?}"
168            );
169            assert_eq!(state.successful_fallback, success, "{reason:?}");
170            if !success {
171                state.apply(run, &committed(&b, ModelFinishReason::ToolCalls));
172                state.apply(run, &committed(&c, reason));
173                state.apply(
174                    run,
175                    &EventPayload::RunFailed {
176                        error: crate::SafeErrorMessage::new("later run failure").unwrap(),
177                        model_error: None,
178                        resolved_model: None,
179                    },
180                );
181                state.apply(run, &EventPayload::RunCancelled { reason: None });
182                assert_eq!(state.selection, Some(b.clone()));
183                assert!(state.successful_fallback);
184                assert_eq!(state.continuation(&a, &initial.chain, false), b);
185            }
186        }
187    }
188
189    #[test]
190    fn continuation_uses_exact_identity_and_respects_scope_and_reset() {
191        let a = selection("a");
192        let b = selection("b");
193        let c = selection("c");
194        let state = SessionModelState {
195            selection: Some(b.clone()),
196            successful_fallback: true,
197            ..Default::default()
198        };
199        let chain = vec![a.model.clone(), b.model.clone(), c.model.clone()];
200        assert_eq!(state.continuation(&a, &chain, false), b);
201        assert_eq!(state.continuation(&a, &chain, true), a);
202        assert_eq!(state.continuation(&c, &chain, false), c);
203        assert_eq!(
204            state.continuation(
205                &a,
206                &[a.model.clone(), c.model.clone(), b.model.clone()],
207                false
208            ),
209            b
210        );
211        assert_eq!(
212            state.continuation(&a, &[a.model.clone(), c.model], false),
213            a
214        );
215        let mut other = a.clone();
216        other.agent = "other".parse().unwrap();
217        assert_eq!(state.continuation(&other, &chain, false), other);
218        other = a.clone();
219        other.preset = Some("python".into());
220        assert_eq!(state.continuation(&other, &chain, false), other);
221        let mut variant_b = b;
222        variant_b.model.variant = Some("fast".parse().unwrap());
223        let state = SessionModelState {
224            selection: Some(variant_b.clone()),
225            successful_fallback: true,
226            ..Default::default()
227        };
228        assert_eq!(state.continuation(&a, &chain, false), a);
229        assert_eq!(
230            state.continuation(&a, &[a.model.clone(), variant_b.model.clone()], false),
231            variant_b
232        );
233    }
234
235    #[test]
236    fn reset_flag_is_optional_and_round_trips_without_changing_events() {
237        let wire = serde_json::json!({"session_id":crate::SessionId::new_v7(),"client_run_id":"test","selection":selection("a"),"input":"hello"});
238        let mut params: crate::RunStartParams = serde_json::from_value(wire.clone()).unwrap();
239        assert!(!params.reset_fallback);
240        assert_eq!(serde_json::to_value(&params).unwrap(), wire);
241        params.reset_fallback = true;
242        assert_eq!(
243            serde_json::to_value(params).unwrap()["reset_fallback"],
244            true
245        );
246    }
247}
cookie_agent_protocol/session_model.rs

cookie_agent_protocol/
session_model.rs
