myco_agent/
state.rs

1//! Deterministic model/tool transitions. Effects carry identities so late or
2//! duplicate completions cannot mutate context or dispatch another tool batch.
3
4use myco_model::{
5    Content, GenerateOutput, Message, TokenUsage, ToolResult, ToolUse, TurnEndReason,
6    answer_content,
7};
8
9use crate::CONTINUE_PROMPT;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
12pub struct OperationId(u64);
13
14/// Durable intent. A tool intent without its result is an unknown external
15/// effect after restart, even when dispatch may not actually have begun.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
17pub enum PendingOperation {
18    Generation { operation: OperationId },
19    Tools { operation: OperationId },
20}
21
22#[derive(Debug, Clone)]
23pub enum Effect {
24    Generate {
25        operation: OperationId,
26    },
27    ExecuteTools {
28        operation: OperationId,
29        calls: Vec<ToolUse>,
30    },
31    Finished {
32        answer: Vec<Content>,
33        reason: TurnEndReason,
34    },
35    Cancelled,
36}
37
38#[derive(Debug, Clone)]
39enum Next {
40    Generate,
41    Finish {
42        answer: Vec<Content>,
43        reason: TurnEndReason,
44    },
45}
46
47#[derive(Debug, Clone)]
48enum Phase {
49    Ready,
50    Finished {
51        answer: Vec<Content>,
52        reason: TurnEndReason,
53    },
54    Cancelled,
55    Generating(OperationId),
56    Tools {
57        operation: OperationId,
58        count: usize,
59        next: Next,
60    },
61}
62
63#[derive(Debug, Clone, thiserror::Error, PartialEq, Eq)]
64pub enum StateError {
65    #[error("agent has an outstanding operation")]
66    Busy,
67    #[error("completion does not match the outstanding operation")]
68    UnexpectedCompletion,
69    #[error("invalid context: {0}")]
70    InvalidContext(String),
71    #[error("invalid model response: {0}")]
72    InvalidResponse(String),
73}
74
75/// Context and control state, without a model, executor, clock, or event sink.
76/// Every transition is synchronous; the caller interprets returned effects.
77#[derive(Debug, Clone)]
78pub struct AgentState {
79    history: Vec<Message>,
80    last_usage: Option<TokenUsage>,
81    phase: Phase,
82    sequence: u64,
83    executing: bool,
84    run_usage: Option<TokenUsage>,
85    truncations: u32,
86    max_truncated_resumes: u32,
87}
88
89impl Default for AgentState {
90    fn default() -> Self {
91        Self {
92            history: Vec::new(),
93            last_usage: None,
94            phase: Phase::Ready,
95            sequence: 0,
96            executing: false,
97            run_usage: None,
98            truncations: 0,
99            max_truncated_resumes: crate::DEFAULT_MAX_TRUNCATED_RESUMES,
100        }
101    }
102}
103
104impl AgentState {
105    pub fn history(&self) -> &[Message] {
106        &self.history
107    }
108
109    pub fn last_usage(&self) -> Option<TokenUsage> {
110        self.last_usage
111    }
112
113    pub fn run_usage(&self) -> Option<TokenUsage> {
114        self.run_usage
115    }
116
117    pub fn is_idle(&self) -> bool {
118        matches!(
119            self.phase,
120            Phase::Ready | Phase::Finished { .. } | Phase::Cancelled
121        )
122    }
123
124    pub fn pending_operation(&self) -> Option<PendingOperation> {
125        match self.phase {
126            Phase::Ready | Phase::Finished { .. } | Phase::Cancelled => None,
127            Phase::Generating(operation) => Some(PendingOperation::Generation { operation }),
128            Phase::Tools { operation, .. } => Some(PendingOperation::Tools { operation }),
129        }
130    }
131
132    /// A pending tool batch leaves a call/result boundary incomplete. It must
133    /// settle before this history can be used as model input or forked.
134    pub fn is_model_boundary(&self) -> bool {
135        !matches!(self.phase, Phase::Tools { .. })
136    }
137
138    /// The next effect, or terminal outcome, without advancing the controller.
139    pub fn effect(&self) -> Option<Effect> {
140        match &self.phase {
141            Phase::Ready => None,
142            Phase::Generating(operation) => Some(Effect::Generate {
143                operation: *operation,
144            }),
145            Phase::Tools { operation, .. } => {
146                let Some(Message::AssistantMessage { tool_uses, .. }) = self.history.last() else {
147                    unreachable!()
148                };
149                Some(Effect::ExecuteTools {
150                    operation: *operation,
151                    calls: tool_uses.clone(),
152                })
153            }
154            Phase::Finished { answer, reason } => Some(Effect::Finished {
155                answer: answer.clone(),
156                reason: reason.clone(),
157            }),
158            Phase::Cancelled => Some(Effect::Cancelled),
159        }
160    }
161
162    /// An interpreter must not dispatch a dropped future's operation again.
163    pub fn begin_effect(&mut self, operation: OperationId) -> Result<(), StateError> {
164        if self.executing {
165            return Err(StateError::Busy);
166        }
167        match self.pending_operation() {
168            Some(
169                PendingOperation::Generation {
170                    operation: expected,
171                }
172                | PendingOperation::Tools {
173                    operation: expected,
174                },
175            ) if expected == operation => {
176                self.executing = true;
177                Ok(())
178            }
179            _ => Err(StateError::UnexpectedCompletion),
180        }
181    }
182
183    pub fn can_replace_at_boundary(&self) -> bool {
184        matches!(self.phase, Phase::Generating(_) | Phase::Finished { .. }) && !self.executing
185    }
186
187    pub fn cancel_at_boundary(&mut self) -> Result<(), StateError> {
188        if !self.can_replace_at_boundary() {
189            return Err(StateError::Busy);
190        }
191        self.phase = Phase::Cancelled;
192        Ok(())
193    }
194
195    /// Install derived context between generations while retaining run policy
196    /// counters. The old operation identity can no longer complete this state.
197    pub fn replace_at_boundary(
198        &mut self,
199        history: Vec<Message>,
200        usage: Option<TokenUsage>,
201    ) -> Result<(), StateError> {
202        if !self.can_replace_at_boundary() {
203            return Err(StateError::Busy);
204        }
205        validate_context(&history)?;
206        self.history = history;
207        self.last_usage = usage;
208        self.generate();
209        Ok(())
210    }
211
212    /// Explicitly abandon work whose future was dropped. Completed observations
213    /// survive; an executing tool batch is reconciled as unknown, never replayed.
214    pub fn recover_interrupted(&mut self) -> Result<(), StateError> {
215        self.history = recover_checkpoint(self.history.clone(), self.pending_operation())?;
216        self.phase = Phase::Ready;
217        self.executing = false;
218        Ok(())
219    }
220
221    fn require_ready(&self) -> Result<(), StateError> {
222        if self.is_idle() {
223            Ok(())
224        } else {
225            Err(StateError::Busy)
226        }
227    }
228
229    pub fn replace_context(
230        &mut self,
231        history: Vec<Message>,
232        usage: Option<TokenUsage>,
233    ) -> Result<(), StateError> {
234        self.require_ready()?;
235        validate_context(&history)?;
236        self.history = history;
237        self.last_usage = usage;
238        self.phase = Phase::Ready;
239        Ok(())
240    }
241
242    pub fn append_input(&mut self, message: Message) -> Result<(), StateError> {
243        self.require_ready()?;
244        if !matches!(message, Message::UserMessage { .. }) {
245            return Err(StateError::InvalidContext(
246                "input must be a user message".into(),
247            ));
248        }
249        self.phase = Phase::Ready;
250        self.history.push(message);
251        Ok(())
252    }
253
254    /// Record runtime observations at a settled boundary without changing the
255    /// run's result, usage, or truncation counters.
256    pub fn append_system(&mut self, parts: Vec<Content>) -> Result<(), StateError> {
257        if self.executing || !self.is_model_boundary() {
258            return Err(StateError::Busy);
259        }
260        if parts.is_empty()
261            || parts
262                .iter()
263                .any(|part| !matches!(part, Content::System { .. }))
264        {
265            return Err(StateError::InvalidContext(
266                "runtime observations must be system parts".into(),
267            ));
268        }
269        self.history.push(Message::UserMessage { content: parts });
270        if matches!(self.phase, Phase::Generating(_)) {
271            self.generate();
272        }
273        Ok(())
274    }
275
276    pub fn truncate_history(&mut self, index: usize) -> Result<Vec<Message>, StateError> {
277        self.require_ready()?;
278        let prefix = self.history.get(..index).ok_or_else(|| {
279            StateError::InvalidContext("rewind index is beyond the history".into())
280        })?;
281        validate_context(prefix)?;
282        let dropped = self.history.split_off(index);
283        self.last_usage = None;
284        self.phase = Phase::Ready;
285        Ok(dropped)
286    }
287
288    pub fn set_max_truncated_resumes(&mut self, resumes: u32) {
289        self.max_truncated_resumes = resumes;
290    }
291
292    pub fn start(&mut self) -> Result<Effect, StateError> {
293        self.require_ready()?;
294        if self.history.is_empty() {
295            return Err(StateError::InvalidContext(
296                "cannot run an empty context".into(),
297            ));
298        }
299        self.run_usage = None;
300        self.truncations = 0;
301        Ok(self.generate())
302    }
303
304    fn operation(&mut self) -> OperationId {
305        self.sequence = self
306            .sequence
307            .checked_add(1)
308            .expect("operation sequence exhausted");
309        OperationId(self.sequence)
310    }
311
312    fn generate(&mut self) -> Effect {
313        let operation = self.operation();
314        self.phase = Phase::Generating(operation);
315        self.executing = false;
316        Effect::Generate { operation }
317    }
318
319    /// Stop a failed or cancelled generation without adding a partial response.
320    pub fn generation_failed(&mut self, operation: OperationId) -> Result<(), StateError> {
321        if !matches!(self.phase, Phase::Generating(id) if id == operation) {
322            return Err(StateError::UnexpectedCompletion);
323        }
324        self.phase = Phase::Ready;
325        self.executing = false;
326        Ok(())
327    }
328
329    pub(crate) fn append_generation_notice(
330        &mut self,
331        operation: OperationId,
332        text: String,
333    ) -> Result<(), StateError> {
334        if !matches!(self.phase, Phase::Generating(id) if id == operation) {
335            return Err(StateError::UnexpectedCompletion);
336        }
337        let notice = Content::System {
338            kind: "generation_notice".into(),
339            text,
340            data: serde_json::Value::Null,
341        };
342        // Notices belong to the current input so user-turn indexes and tool
343        // call/result pairing remain stable across rewind and compaction.
344        match self.history.last_mut() {
345            Some(Message::UserMessage { content }) => content.push(notice),
346            Some(Message::ToolResults { tool_use_results }) if !tool_use_results.is_empty() => {
347                tool_use_results.last_mut().unwrap().content.push(notice);
348            }
349            _ => self.history.push(Message::UserMessage {
350                content: vec![notice],
351            }),
352        }
353        Ok(())
354    }
355
356    pub fn generated(
357        &mut self,
358        operation: OperationId,
359        output: GenerateOutput,
360    ) -> Result<Effect, StateError> {
361        if !matches!(self.phase, Phase::Generating(id) if id == operation) {
362            return Err(StateError::UnexpectedCompletion);
363        }
364        self.executing = false;
365        if output
366            .content
367            .iter()
368            .any(|part| matches!(part, Content::System { .. }))
369        {
370            self.phase = Phase::Ready;
371            return Err(StateError::InvalidResponse(
372                "model output contains a runtime-only system part".into(),
373            ));
374        }
375        if let Some(usage) = output.usage {
376            self.run_usage = Some(TokenUsage {
377                output_tokens: self
378                    .run_usage
379                    .map_or(0, |previous| previous.output_tokens)
380                    .saturating_add(usage.output_tokens),
381                ..usage
382            });
383            self.last_usage = self.run_usage;
384        }
385        let answer = answer_content(&output.content);
386        let reason = output.turn_end_reason;
387        let calls = output.tool_uses;
388        self.history.push(Message::AssistantMessage {
389            content: output.content,
390            tool_uses: calls.clone(),
391            turn_end_reason: Some(reason.clone()),
392        });
393        self.phase = Phase::Ready;
394        if reason == TurnEndReason::ToolUse && calls.is_empty() {
395            return Err(StateError::InvalidResponse(
396                "turn ended in tool_use but streamed zero tool uses".into(),
397            ));
398        }
399        let resume = if reason == TurnEndReason::MaxTokens {
400            self.truncations = self.truncations.saturating_add(1);
401            self.truncations <= self.max_truncated_resumes
402        } else {
403            self.truncations = 0;
404            false
405        };
406        let next = if reason == TurnEndReason::ToolUse || resume {
407            Next::Generate
408        } else {
409            Next::Finish { answer, reason }
410        };
411        if calls.is_empty() {
412            if resume {
413                self.history.push(Message::UserMessage {
414                    content: vec![Content::System {
415                        kind: "continuation".into(),
416                        text: CONTINUE_PROMPT.into(),
417                        data: serde_json::json!({"reason":"max_tokens"}),
418                    }],
419                });
420            }
421            Ok(self.advance(next))
422        } else {
423            let operation = self.operation();
424            self.phase = Phase::Tools {
425                operation,
426                count: calls.len(),
427                next,
428            };
429            Ok(Effect::ExecuteTools { operation, calls })
430        }
431    }
432
433    /// Results are in call order, including cancellation or unknown outcomes.
434    /// Cancellation settles the entire batch and prevents another generation.
435    pub fn tools_completed(
436        &mut self,
437        operation: OperationId,
438        results: Vec<ToolResult>,
439        cancelled: bool,
440    ) -> Result<Effect, StateError> {
441        let Phase::Tools {
442            operation: expected,
443            count,
444            next,
445        } = &self.phase
446        else {
447            return Err(StateError::UnexpectedCompletion);
448        };
449        if operation != *expected {
450            return Err(StateError::UnexpectedCompletion);
451        }
452        if results.len() != *count {
453            return Err(StateError::InvalidContext(format!(
454                "expected {count} tool results, received {}",
455                results.len(),
456            )));
457        }
458        let next = next.clone();
459        self.executing = false;
460        self.history.push(Message::ToolResults {
461            tool_use_results: results,
462        });
463        if cancelled {
464            self.phase = Phase::Cancelled;
465            Ok(Effect::Cancelled)
466        } else {
467            Ok(self.advance(next))
468        }
469    }
470
471    fn advance(&mut self, next: Next) -> Effect {
472        match next {
473            Next::Generate => self.generate(),
474            Next::Finish { answer, reason } => {
475                self.phase = Phase::Finished {
476                    answer: answer.clone(),
477                    reason: reason.clone(),
478                };
479                Effect::Finished { answer, reason }
480            }
481        }
482    }
483}
484
485/// Every tool call has exactly one result in the immediately following batch.
486pub fn validate_context(history: &[Message]) -> Result<(), StateError> {
487    let mut pending = 0;
488    for (index, message) in history.iter().enumerate() {
489        match message {
490            Message::ToolResults { tool_use_results } => {
491                if pending == 0 || pending != tool_use_results.len() {
492                    return Err(StateError::InvalidContext(format!(
493                        "message {index} has {} results for {pending} tool calls",
494                        tool_use_results.len(),
495                    )));
496                }
497                pending = 0;
498            }
499            _ if pending != 0 => {
500                return Err(StateError::InvalidContext(format!(
501                    "message {index} interrupts an unanswered tool batch",
502                )));
503            }
504            Message::AssistantMessage { tool_uses, .. } => pending = tool_uses.len(),
505            Message::UserMessage { .. } => {}
506        }
507    }
508    if pending != 0 {
509        return Err(StateError::InvalidContext(
510            "history ends with unanswered tool calls".into(),
511        ));
512    }
513    Ok(())
514}
515
516pub fn validate_checkpoint(
517    history: &[Message],
518    pending: Option<PendingOperation>,
519) -> Result<(), StateError> {
520    if matches!(pending, Some(PendingOperation::Tools { .. })) {
521        match history.split_last() {
522            Some((Message::AssistantMessage { tool_uses, .. }, prefix))
523                if !tool_uses.is_empty() =>
524            {
525                validate_context(prefix)
526            }
527            _ => Err(StateError::InvalidContext(
528                "pending tools require an unanswered assistant tool batch".into(),
529            )),
530        }
531    } else {
532        validate_context(history)
533    }
534}
535
536/// Recover observations without reissuing any pending operation. The caller
537/// must persist this context before allowing new effects.
538pub fn recover_checkpoint(
539    mut history: Vec<Message>,
540    pending: Option<PendingOperation>,
541) -> Result<Vec<Message>, StateError> {
542    validate_checkpoint(&history, pending)?;
543    let Some(pending) = pending else {
544        return Ok(history);
545    };
546    let text = match pending {
547        PendingOperation::Tools { .. } => {
548            let Some(Message::AssistantMessage { tool_uses, .. }) = history.last() else {
549                unreachable!()
550            };
551            let results = tool_uses.iter().map(|_| ToolResult::err(
552                "execution interrupted before the outcome was saved; effects are unknown. Inspect external state before retrying this action.",
553            )).collect();
554            history.push(Message::ToolResults {
555                tool_use_results: results,
556            });
557            "This runtime recovered an interrupted tool batch. Its results were not durably recorded. Calls may have executed; they have not been replayed. Inspect external state before repeating any action."
558        }
559        PendingOperation::Generation { .. } => {
560            "A generation was pending when the previous runtime stopped. No completed response from that operation was recorded. Continue from the saved observations."
561        }
562    };
563    history.push(Message::UserMessage {
564        content: vec![Content::System {
565            kind: "interrupted".into(),
566            text: text.into(),
567            data: serde_json::json!({"pending":pending}),
568        }],
569    });
570    Ok(history)
571}
572
573#[cfg(test)]
574mod tests {
575    use super::*;
576    use crate::test_support::{assistant, assistant_tool, tool_results, user};
577    use serde_json::json;
578
579    fn state() -> AgentState {
580        let mut state = AgentState::default();
581        state.append_input(user("task")).unwrap();
582        state
583    }
584
585    fn generation(effect: Effect) -> OperationId {
586        match effect {
587            Effect::Generate { operation } => operation,
588            other => panic!("expected generation, got {other:?}"),
589        }
590    }
591
592    fn response(reason: TurnEndReason, calls: usize) -> GenerateOutput {
593        GenerateOutput {
594            content: vec![Content::Text {
595                text: "answer".into(),
596            }],
597            tool_uses: (0..calls)
598                .map(|index| ToolUse {
599                    name: "test".into(),
600                    input: json!({"index": index}),
601                })
602                .collect(),
603            turn_end_reason: reason,
604            usage: Some(TokenUsage {
605                input_tokens: 100,
606                output_tokens: 10,
607                cached_input_tokens: 20,
608            }),
609        }
610    }
611
612    #[test]
613    fn pure_transitions_replay_the_same_tool_loop_without_a_runtime() {
614        let mut left = state();
615        let mut right = left.clone();
616        for machine in [&mut left, &mut right] {
617            let op = generation(machine.start().unwrap());
618            let Effect::ExecuteTools { operation, calls } = machine
619                .generated(op, response(TurnEndReason::ToolUse, 2))
620                .unwrap()
621            else {
622                panic!()
623            };
624            assert_eq!(calls.len(), 2);
625            assert!(!machine.is_model_boundary());
626            let next = machine
627                .tools_completed(
628                    operation,
629                    vec![ToolResult::text("first"), ToolResult::err("second")],
630                    false,
631                )
632                .unwrap();
633            let op = generation(next);
634            assert!(matches!(
635                machine
636                    .generated(op, response(TurnEndReason::EndTurn, 0))
637                    .unwrap(),
638                Effect::Finished {
639                    reason: TurnEndReason::EndTurn,
640                    ..
641                }
642            ));
643            assert_eq!(machine.last_usage().unwrap().output_tokens, 20);
644            validate_context(machine.history()).unwrap();
645        }
646        assert_eq!(format!("{left:?}"), format!("{right:?}"));
647    }
648
649    #[test]
650    fn stale_completions_and_input_during_work_leave_state_unchanged() {
651        let mut machine = state();
652        let stale = generation(machine.start().unwrap());
653        machine.generation_failed(stale).unwrap();
654        let current = generation(machine.start().unwrap());
655        let before = format!("{machine:?}");
656        assert_eq!(
657            machine
658                .generated(stale, response(TurnEndReason::ToolUse, 1))
659                .unwrap_err(),
660            StateError::UnexpectedCompletion
661        );
662        assert_eq!(
663            machine.append_input(user("interjection")),
664            Err(StateError::Busy)
665        );
666        assert_eq!(
667            machine.replace_context(vec![user("replacement")], None),
668            Err(StateError::Busy)
669        );
670        assert!(matches!(machine.truncate_history(0), Err(StateError::Busy)));
671        assert_eq!(
672            machine.generation_failed(stale),
673            Err(StateError::UnexpectedCompletion)
674        );
675        assert_eq!(format!("{machine:?}"), before);
676        machine.generation_failed(current).unwrap();
677    }
678
679    #[test]
680    fn tool_batch_requires_matching_count_and_can_only_complete_once() {
681        let mut machine = state();
682        let generation = generation(machine.start().unwrap());
683        let Effect::ExecuteTools { operation, .. } = machine
684            .generated(generation, response(TurnEndReason::ToolUse, 2))
685            .unwrap()
686        else {
687            panic!()
688        };
689        let before = format!("{machine:?}");
690        assert!(
691            machine
692                .tools_completed(operation, vec![ToolResult::text("partial")], false)
693                .is_err()
694        );
695        assert_eq!(format!("{machine:?}"), before);
696        assert!(matches!(
697            machine
698                .tools_completed(
699                    operation,
700                    vec![ToolResult::text("finished"), ToolResult::err("cancelled")],
701                    true
702                )
703                .unwrap(),
704            Effect::Cancelled
705        ));
706        validate_context(machine.history()).unwrap();
707        assert_eq!(
708            machine
709                .tools_completed(operation, vec![], false)
710                .unwrap_err(),
711            StateError::UnexpectedCompletion
712        );
713        machine.append_input(user("next turn")).unwrap();
714    }
715
716    #[test]
717    fn context_replacement_and_rewind_reject_split_or_orphaned_tool_batches() {
718        let mut machine = state();
719        let call = assistant_tool(None, "test", json!({}));
720        for history in [
721            vec![call.clone()],
722            vec![tool_results(&["orphan"])],
723            vec![call.clone(), user("interrupt"), tool_results(&["result"])],
724            vec![call.clone(), tool_results(&["extra", "extra"])],
725            vec![Message::ToolResults {
726                tool_use_results: vec![],
727            }],
728        ] {
729            assert!(machine.replace_context(history, None).is_err());
730            assert_eq!(machine.history().len(), 1);
731        }
732        machine
733            .replace_context(
734                vec![user("task"), call, tool_results(&["ok"]), assistant("done")],
735                None,
736            )
737            .unwrap();
738        assert!(machine.truncate_history(2).is_err());
739        assert!(machine.truncate_history(99).is_err());
740        assert_eq!(machine.history().len(), 4);
741        machine.truncate_history(1).unwrap();
742    }
743
744    #[test]
745    fn truncation_cap_is_an_explicit_finish_reason_and_resets_for_a_new_run() {
746        let mut machine = state();
747        machine.set_max_truncated_resumes(1);
748        let first = generation(machine.start().unwrap());
749        let next = machine
750            .generated(first, response(TurnEndReason::MaxTokens, 0))
751            .unwrap();
752        let second = generation(next);
753        assert!(matches!(
754            machine
755                .generated(second, response(TurnEndReason::MaxTokens, 0))
756                .unwrap(),
757            Effect::Finished {
758                reason: TurnEndReason::MaxTokens,
759                ..
760            }
761        ));
762        machine.append_input(user("new task")).unwrap();
763        let next = generation(machine.start().unwrap());
764        assert!(matches!(
765            machine
766                .generated(next, response(TurnEndReason::MaxTokens, 0))
767                .unwrap(),
768            Effect::Generate { .. }
769        ));
770    }
771}
myco_agent/state.rs

myco_agent/
state.rs
