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, LOST_THINKING_PROMPT, lost_thinking};
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    GenerationFailed,
51    Finished {
52        answer: Vec<Content>,
53        reason: TurnEndReason,
54    },
55    Cancelled,
56    Generating(OperationId),
57    Tools {
58        operation: OperationId,
59        count: usize,
60        next: Next,
61    },
62}
63
64#[derive(Debug, Clone, thiserror::Error, PartialEq, Eq)]
65pub enum StateError {
66    #[error("agent has an outstanding operation")]
67    Busy,
68    #[error("completion does not match the outstanding operation")]
69    UnexpectedCompletion,
70    #[error("invalid context: {0}")]
71    InvalidContext(String),
72    #[error("invalid model response: {0}")]
73    InvalidResponse(String),
74}
75
76/// Context and control state, without a model, executor, clock, or event sink.
77/// Every transition is synchronous; the caller interprets returned effects.
78#[derive(Debug, Clone)]
79pub struct AgentState {
80    history: Vec<Message>,
81    context_size: crate::context_size::ContextSize,
82    last_usage: Option<TokenUsage>,
83    phase: Phase,
84    sequence: u64,
85    executing: bool,
86    run_usage: Option<TokenUsage>,
87    truncations: u32,
88    max_truncated_resumes: u32,
89}
90
91impl Default for AgentState {
92    fn default() -> Self {
93        Self {
94            history: Vec::new(),
95            context_size: Default::default(),
96            last_usage: None,
97            phase: Phase::Ready,
98            sequence: 0,
99            executing: false,
100            run_usage: None,
101            truncations: 0,
102            max_truncated_resumes: crate::DEFAULT_MAX_TRUNCATED_RESUMES,
103        }
104    }
105}
106
107impl AgentState {
108    pub fn history(&self) -> &[Message] {
109        &self.history
110    }
111
112    pub fn last_usage(&self) -> Option<TokenUsage> {
113        self.last_usage
114    }
115
116    pub(crate) fn context_tokens_estimate(&self) -> u64 {
117        self.context_size.tokens()
118    }
119
120    pub(crate) fn restore_context_size_hint(&mut self, tokens: Option<u64>) {
121        self.context_size.restore_hint(tokens);
122    }
123
124    fn push(&mut self, message: Message) {
125        self.context_size.append(&message);
126        self.history.push(message);
127    }
128
129    fn replace_history(&mut self, history: Vec<Message>, usage: Option<TokenUsage>) {
130        if (usage.is_none() || usage == self.last_usage) && history.starts_with(&self.history) {
131            for message in &history[self.history.len()..] {
132                self.context_size.append(message);
133            }
134        } else {
135            self.context_size = crate::context_size::ContextSize::restored(&history, usage);
136        }
137        self.history = history;
138        self.last_usage = usage;
139    }
140
141    pub fn run_usage(&self) -> Option<TokenUsage> {
142        self.run_usage
143    }
144
145    pub fn is_idle(&self) -> bool {
146        matches!(
147            self.phase,
148            Phase::Ready | Phase::GenerationFailed | Phase::Finished { .. } | Phase::Cancelled
149        )
150    }
151
152    pub fn pending_operation(&self) -> Option<PendingOperation> {
153        match self.phase {
154            Phase::Ready | Phase::GenerationFailed | Phase::Finished { .. } | Phase::Cancelled => {
155                None
156            }
157            Phase::Generating(operation) => Some(PendingOperation::Generation { operation }),
158            Phase::Tools { operation, .. } => Some(PendingOperation::Tools { operation }),
159        }
160    }
161
162    /// A pending tool batch leaves a call/result boundary incomplete. It must
163    /// settle before this history can be used as model input or forked.
164    pub fn is_model_boundary(&self) -> bool {
165        !matches!(self.phase, Phase::Tools { .. })
166    }
167
168    /// The next effect, or terminal outcome, without advancing the controller.
169    pub fn effect(&self) -> Option<Effect> {
170        match &self.phase {
171            Phase::Ready | Phase::GenerationFailed => None,
172            Phase::Generating(operation) => Some(Effect::Generate {
173                operation: *operation,
174            }),
175            Phase::Tools { operation, .. } => {
176                let Some(Message::AssistantMessage { tool_uses, .. }) = self.history.last() else {
177                    unreachable!()
178                };
179                Some(Effect::ExecuteTools {
180                    operation: *operation,
181                    calls: tool_uses.clone(),
182                })
183            }
184            Phase::Finished { answer, reason } => Some(Effect::Finished {
185                answer: answer.clone(),
186                reason: reason.clone(),
187            }),
188            Phase::Cancelled => Some(Effect::Cancelled),
189        }
190    }
191
192    /// An interpreter must not dispatch a dropped future's operation again.
193    pub fn begin_effect(&mut self, operation: OperationId) -> Result<(), StateError> {
194        if self.executing {
195            return Err(StateError::Busy);
196        }
197        match self.pending_operation() {
198            Some(
199                PendingOperation::Generation {
200                    operation: expected,
201                }
202                | PendingOperation::Tools {
203                    operation: expected,
204                },
205            ) if expected == operation => {
206                self.executing = true;
207                Ok(())
208            }
209            _ => Err(StateError::UnexpectedCompletion),
210        }
211    }
212
213    pub fn can_replace_at_boundary(&self) -> bool {
214        matches!(
215            self.phase,
216            Phase::Generating(_) | Phase::GenerationFailed | Phase::Finished { .. }
217        ) && !self.executing
218    }
219
220    pub fn cancel_at_boundary(&mut self) -> Result<(), StateError> {
221        if !self.can_replace_at_boundary() {
222            return Err(StateError::Busy);
223        }
224        self.phase = Phase::Cancelled;
225        Ok(())
226    }
227
228    /// Install derived context between generations while retaining run policy
229    /// counters. The old operation identity can no longer complete this state.
230    pub fn replace_at_boundary(
231        &mut self,
232        history: Vec<Message>,
233        usage: Option<TokenUsage>,
234    ) -> Result<(), StateError> {
235        if !self.can_replace_at_boundary() {
236            return Err(StateError::Busy);
237        }
238        validate_context(&history)?;
239        self.replace_history(history, usage);
240        self.generate();
241        Ok(())
242    }
243
244    /// Explicitly abandon work whose future was dropped. Completed observations
245    /// survive; an executing tool batch is reconciled as unknown, never replayed.
246    /// Live intents that never began have known, unexecuted outcomes.
247    pub fn recover_interrupted(&mut self) -> Result<(), StateError> {
248        if self.executing {
249            let history = recover_checkpoint(self.history.clone(), self.pending_operation())?;
250            self.replace_history(history, self.last_usage);
251        } else if let Phase::Tools {
252            operation, count, ..
253        } = self.phase
254        {
255            let results = (0..count).map(|_| ToolResult::err(
256                "not executed: the run stopped before this tool batch began. Reconsider this action using the latest input before trying again.",
257            )).collect();
258            self.tools_completed(operation, results, true)?;
259        }
260        self.phase = Phase::Ready;
261        self.executing = false;
262        Ok(())
263    }
264
265    fn require_ready(&self) -> Result<(), StateError> {
266        if self.is_idle() {
267            Ok(())
268        } else {
269            Err(StateError::Busy)
270        }
271    }
272
273    pub fn replace_context(
274        &mut self,
275        history: Vec<Message>,
276        usage: Option<TokenUsage>,
277    ) -> Result<(), StateError> {
278        self.require_ready()?;
279        validate_context(&history)?;
280        self.replace_history(history, usage);
281        self.phase = Phase::Ready;
282        Ok(())
283    }
284
285    pub fn append_input(&mut self, message: Message) -> Result<(), StateError> {
286        self.require_ready()?;
287        if !matches!(message, Message::UserMessage { .. }) {
288            return Err(StateError::InvalidContext(
289                "input must be a user message".into(),
290            ));
291        }
292        self.phase = Phase::Ready;
293        self.push(message);
294        Ok(())
295    }
296
297    /// Record runtime observations at a settled boundary without changing the
298    /// run's result, usage, or truncation counters.
299    pub fn append_system(&mut self, parts: Vec<Content>) -> Result<(), StateError> {
300        if self.executing || !self.is_model_boundary() {
301            return Err(StateError::Busy);
302        }
303        if parts.is_empty()
304            || parts
305                .iter()
306                .any(|part| !matches!(part, Content::System { .. }))
307        {
308            return Err(StateError::InvalidContext(
309                "runtime observations must be system parts".into(),
310            ));
311        }
312        self.push(Message::UserMessage { content: parts });
313        if matches!(self.phase, Phase::Generating(_)) {
314            self.generate();
315        }
316        Ok(())
317    }
318
319    pub fn truncate_history(&mut self, index: usize) -> Result<Vec<Message>, StateError> {
320        self.require_ready()?;
321        let prefix = self.history.get(..index).ok_or_else(|| {
322            StateError::InvalidContext("rewind index is beyond the history".into())
323        })?;
324        validate_context(prefix)?;
325        let dropped = self.history.split_off(index);
326        self.context_size = crate::context_size::ContextSize::restored(&self.history, None);
327        self.last_usage = None;
328        self.phase = Phase::Ready;
329        Ok(dropped)
330    }
331
332    pub fn set_max_truncated_resumes(&mut self, resumes: u32) {
333        self.max_truncated_resumes = resumes;
334    }
335
336    /// Retry a failed generation within its existing run budget. Appending new
337    /// input first transitions to Ready and begins a fresh run instead.
338    pub fn start(&mut self) -> Result<Effect, StateError> {
339        self.require_ready()?;
340        if self.history.is_empty() {
341            return Err(StateError::InvalidContext(
342                "cannot run an empty context".into(),
343            ));
344        }
345        if !matches!(self.phase, Phase::GenerationFailed) {
346            self.run_usage = None;
347            self.truncations = 0;
348        }
349        Ok(self.generate())
350    }
351
352    fn operation(&mut self) -> OperationId {
353        self.sequence = self
354            .sequence
355            .checked_add(1)
356            .expect("operation sequence exhausted");
357        OperationId(self.sequence)
358    }
359
360    fn generate(&mut self) -> Effect {
361        let operation = self.operation();
362        self.phase = Phase::Generating(operation);
363        self.executing = false;
364        Effect::Generate { operation }
365    }
366
367    /// Stop a failed or cancelled generation without adding a partial response.
368    /// Context can be repaired at this boundary without resetting run counters.
369    pub fn generation_failed(&mut self, operation: OperationId) -> Result<(), StateError> {
370        if !matches!(self.phase, Phase::Generating(id) if id == operation) {
371            return Err(StateError::UnexpectedCompletion);
372        }
373        self.phase = Phase::GenerationFailed;
374        self.executing = false;
375        Ok(())
376    }
377
378    pub(crate) fn append_generation_notice(
379        &mut self,
380        operation: OperationId,
381        text: String,
382    ) -> Result<(), StateError> {
383        if !matches!(self.phase, Phase::Generating(id) if id == operation) {
384            return Err(StateError::UnexpectedCompletion);
385        }
386        let notice = Content::System {
387            kind: "generation_notice".into(),
388            text,
389            data: serde_json::Value::Null,
390        };
391        // Notices belong to the current input so user-turn indexes and tool
392        // call/result pairing remain stable across rewind and compaction.
393        match self.history.last_mut() {
394            Some(Message::UserMessage { content }) => {
395                self.context_size.append_part(&notice);
396                content.push(notice);
397            }
398            Some(Message::ToolResults { tool_use_results }) if !tool_use_results.is_empty() => {
399                self.context_size.append_part(&notice);
400                tool_use_results.last_mut().unwrap().content.push(notice);
401            }
402            _ => self.push(Message::UserMessage {
403                content: vec![notice],
404            }),
405        }
406        Ok(())
407    }
408
409    pub fn generated(
410        &mut self,
411        operation: OperationId,
412        output: GenerateOutput,
413    ) -> Result<Effect, StateError> {
414        if !matches!(self.phase, Phase::Generating(id) if id == operation) {
415            return Err(StateError::UnexpectedCompletion);
416        }
417        if output
418            .content
419            .iter()
420            .any(|part| matches!(part, Content::System { .. }))
421        {
422            self.generation_failed(operation)?;
423            return Err(StateError::InvalidResponse(
424                "model output contains a runtime-only system part".into(),
425            ));
426        }
427        if output.turn_end_reason == TurnEndReason::ToolUse && output.tool_uses.is_empty() {
428            self.generation_failed(operation)?;
429            return Err(StateError::InvalidResponse(
430                "turn ended in tool_use but streamed zero tool uses".into(),
431            ));
432        }
433        self.executing = false;
434        if let Some(usage) = output.usage {
435            self.context_size.observe_input(usage);
436            self.run_usage = Some(TokenUsage {
437                output_tokens: self
438                    .run_usage
439                    .map_or(0, |previous| previous.output_tokens)
440                    .saturating_add(usage.output_tokens),
441                ..usage
442            });
443            self.last_usage = self.run_usage;
444        }
445        let answer = answer_content(&output.content);
446        let reason = output.turn_end_reason;
447        let calls = output.tool_uses;
448        let assistant = Message::AssistantMessage {
449            content: output.content,
450            tool_uses: calls.clone(),
451            turn_end_reason: Some(reason.clone()),
452        };
453        let thinking_only = lost_thinking(&assistant);
454        self.push(assistant);
455        self.phase = Phase::Ready;
456        let resume = if reason == TurnEndReason::MaxTokens {
457            self.truncations = self.truncations.saturating_add(1);
458            self.truncations <= self.max_truncated_resumes
459        } else {
460            self.truncations = 0;
461            false
462        };
463        let next = if reason == TurnEndReason::ToolUse || resume {
464            Next::Generate
465        } else {
466            Next::Finish { answer, reason }
467        };
468        if calls.is_empty() {
469            if resume {
470                self.push(Message::UserMessage {
471                    content: vec![Content::System {
472                        kind: "continuation".into(),
473                        text: if thinking_only {
474                            LOST_THINKING_PROMPT
475                        } else {
476                            CONTINUE_PROMPT
477                        }
478                        .into(),
479                        data: serde_json::json!({"reason":"max_tokens","thinking_only":thinking_only}),
480                    }],
481                });
482            }
483            Ok(self.advance(next))
484        } else {
485            let operation = self.operation();
486            self.phase = Phase::Tools {
487                operation,
488                count: calls.len(),
489                next,
490            };
491            Ok(Effect::ExecuteTools { operation, calls })
492        }
493    }
494
495    /// Results are in call order, including cancellation or unknown outcomes.
496    /// Cancellation settles the entire batch and prevents another generation.
497    pub fn tools_completed(
498        &mut self,
499        operation: OperationId,
500        results: Vec<ToolResult>,
501        cancelled: bool,
502    ) -> Result<Effect, StateError> {
503        let Phase::Tools {
504            operation: expected,
505            count,
506            next,
507        } = &self.phase
508        else {
509            return Err(StateError::UnexpectedCompletion);
510        };
511        if operation != *expected {
512            return Err(StateError::UnexpectedCompletion);
513        }
514        if results.len() != *count {
515            return Err(StateError::InvalidContext(format!(
516                "expected {count} tool results, received {}",
517                results.len(),
518            )));
519        }
520        let next = next.clone();
521        self.executing = false;
522        self.push(Message::ToolResults {
523            tool_use_results: results,
524        });
525        if cancelled {
526            self.phase = Phase::Cancelled;
527            Ok(Effect::Cancelled)
528        } else {
529            Ok(self.advance(next))
530        }
531    }
532
533    fn advance(&mut self, next: Next) -> Effect {
534        match next {
535            Next::Generate => self.generate(),
536            Next::Finish { answer, reason } => {
537                self.phase = Phase::Finished {
538                    answer: answer.clone(),
539                    reason: reason.clone(),
540                };
541                Effect::Finished { answer, reason }
542            }
543        }
544    }
545}
546
547/// Every tool call has exactly one result in the immediately following batch.
548pub fn validate_context(history: &[Message]) -> Result<(), StateError> {
549    let mut pending = 0;
550    for (index, message) in history.iter().enumerate() {
551        match message {
552            Message::ToolResults { tool_use_results } => {
553                if pending == 0 || pending != tool_use_results.len() {
554                    return Err(StateError::InvalidContext(format!(
555                        "message {index} has {} results for {pending} tool calls",
556                        tool_use_results.len(),
557                    )));
558                }
559                pending = 0;
560            }
561            _ if pending != 0 => {
562                return Err(StateError::InvalidContext(format!(
563                    "message {index} interrupts an unanswered tool batch",
564                )));
565            }
566            Message::AssistantMessage { tool_uses, .. } => pending = tool_uses.len(),
567            Message::UserMessage { .. } => {}
568        }
569    }
570    if pending != 0 {
571        return Err(StateError::InvalidContext(
572            "history ends with unanswered tool calls".into(),
573        ));
574    }
575    Ok(())
576}
577
578pub fn validate_checkpoint(
579    history: &[Message],
580    pending: Option<PendingOperation>,
581) -> Result<(), StateError> {
582    if matches!(pending, Some(PendingOperation::Tools { .. })) {
583        match history.split_last() {
584            Some((Message::AssistantMessage { tool_uses, .. }, prefix))
585                if !tool_uses.is_empty() =>
586            {
587                validate_context(prefix)
588            }
589            _ => Err(StateError::InvalidContext(
590                "pending tools require an unanswered assistant tool batch".into(),
591            )),
592        }
593    } else {
594        validate_context(history)
595    }
596}
597
598/// Recover observations without reissuing any pending operation. The caller
599/// must persist this context before allowing new effects.
600pub fn recover_checkpoint(
601    mut history: Vec<Message>,
602    pending: Option<PendingOperation>,
603) -> Result<Vec<Message>, StateError> {
604    validate_checkpoint(&history, pending)?;
605    let Some(pending) = pending else {
606        return Ok(history);
607    };
608    let text = match pending {
609        PendingOperation::Tools { .. } => {
610            let Some(Message::AssistantMessage { tool_uses, .. }) = history.last() else {
611                unreachable!()
612            };
613            let results = tool_uses.iter().map(|_| ToolResult::err(
614                "execution interrupted before the outcome was saved; effects are unknown. Inspect external state before retrying this action.",
615            )).collect();
616            history.push(Message::ToolResults {
617                tool_use_results: results,
618            });
619            "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."
620        }
621        PendingOperation::Generation { .. } => {
622            "A generation was pending when the previous runtime stopped. No completed response from that operation was recorded. Continue from the saved observations."
623        }
624    };
625    history.push(Message::UserMessage {
626        content: vec![Content::System {
627            kind: "interrupted".into(),
628            text: text.into(),
629            data: serde_json::json!({"pending":pending}),
630        }],
631    });
632    Ok(history)
633}
634
635#[cfg(test)]
636mod tests {
637    use super::*;
638    use crate::test_support::{assistant, assistant_tool, tool_results, user};
639    use serde_json::json;
640
641    fn state() -> AgentState {
642        let mut state = AgentState::default();
643        state.append_input(user("task")).unwrap();
644        state
645    }
646
647    fn generation(effect: Effect) -> OperationId {
648        match effect {
649            Effect::Generate { operation } => operation,
650            other => panic!("expected generation, got {other:?}"),
651        }
652    }
653
654    fn response(reason: TurnEndReason, calls: usize) -> GenerateOutput {
655        GenerateOutput {
656            content: vec![Content::Text {
657                text: "answer".into(),
658            }],
659            tool_uses: (0..calls)
660                .map(|index| ToolUse {
661                    name: "test".into(),
662                    input: json!({"index": index}),
663                    thought_signature: None,
664                })
665                .collect(),
666            turn_end_reason: reason,
667            usage: Some(TokenUsage {
668                input_tokens: 100,
669                output_tokens: 10,
670                cached_input_tokens: 20,
671            }),
672        }
673    }
674
675    fn continuation_text(machine: &AgentState) -> &str {
676        match machine.history().last() {
677            Some(Message::UserMessage { content }) => match content.as_slice() {
678                [Content::System { text, .. }] => text,
679                other => panic!("expected one continuation part, got {other:?}"),
680            },
681            other => panic!("expected a continuation turn, got {other:?}"),
682        }
683    }
684
685    /// A thinking-only truncation reaches the provider as nothing, so its
686    /// resume prompt must say the thinking was lost instead of "continue".
687    #[test]
688    fn thinking_only_truncation_resumes_with_the_lost_thinking_prompt() {
689        let mut machine = state();
690        let first = generation(machine.start().unwrap());
691        let mut thinking_only = response(TurnEndReason::MaxTokens, 0);
692        thinking_only.content = vec![Content::Thinking {
693            text: "planning".into(),
694            signature: Some("sig".into()),
695            redacted: false,
696        }];
697        let next = generation(machine.generated(first, thinking_only).unwrap());
698        assert_eq!(continuation_text(&machine), LOST_THINKING_PROMPT);
699
700        let mut partial = response(TurnEndReason::MaxTokens, 0);
701        partial.content.insert(
702            0,
703            Content::Thinking {
704                text: "planning".into(),
705                signature: None,
706                redacted: false,
707            },
708        );
709        generation(machine.generated(next, partial).unwrap());
710        assert_eq!(continuation_text(&machine), CONTINUE_PROMPT);
711    }
712
713    #[test]
714    fn pure_transitions_replay_the_same_tool_loop_without_a_runtime() {
715        let mut left = state();
716        let mut right = left.clone();
717        for machine in [&mut left, &mut right] {
718            let op = generation(machine.start().unwrap());
719            let Effect::ExecuteTools { operation, calls } = machine
720                .generated(op, response(TurnEndReason::ToolUse, 2))
721                .unwrap()
722            else {
723                panic!()
724            };
725            assert_eq!(calls.len(), 2);
726            assert!(!machine.is_model_boundary());
727            let next = machine
728                .tools_completed(
729                    operation,
730                    vec![ToolResult::text("first"), ToolResult::err("second")],
731                    false,
732                )
733                .unwrap();
734            let op = generation(next);
735            assert!(matches!(
736                machine
737                    .generated(op, response(TurnEndReason::EndTurn, 0))
738                    .unwrap(),
739                Effect::Finished {
740                    reason: TurnEndReason::EndTurn,
741                    ..
742                }
743            ));
744            assert_eq!(machine.last_usage().unwrap().output_tokens, 20);
745            validate_context(machine.history()).unwrap();
746        }
747        assert_eq!(format!("{left:?}"), format!("{right:?}"));
748    }
749
750    #[test]
751    fn context_growth_survives_rebinding_missing_usage_and_model_changes() {
752        let mut machine = state();
753        let op = generation(machine.start().unwrap());
754        machine
755            .generated(op, response(TurnEndReason::EndTurn, 0))
756            .unwrap();
757        let measured = machine.context_tokens_estimate();
758        machine
759            .replace_context(machine.history().to_vec(), machine.last_usage())
760            .unwrap();
761        assert_eq!(machine.context_tokens_estimate(), measured);
762        machine
763            .append_input(Message::UserMessage {
764                content: vec![Content::Text {
765                    text: "x".repeat(3000),
766                }],
767            })
768            .unwrap();
769        let grown = machine.context_tokens_estimate();
770        assert!(grown >= measured + 1000);
771        let op = generation(machine.start().unwrap());
772        let mut unmeasured = response(TurnEndReason::EndTurn, 0);
773        unmeasured.usage = None;
774        machine.generated(op, unmeasured).unwrap();
775        assert!(machine.context_tokens_estimate() > grown);
776        let before_change = machine.context_tokens_estimate();
777        machine
778            .replace_context(machine.history().to_vec(), None)
779            .unwrap();
780        assert_eq!(machine.context_tokens_estimate(), before_change);
781        machine.truncate_history(1).unwrap();
782        assert!(machine.context_tokens_estimate() < measured);
783    }
784
785    #[test]
786    fn stale_completions_and_input_during_work_leave_state_unchanged() {
787        let mut machine = state();
788        let stale = generation(machine.start().unwrap());
789        machine.generation_failed(stale).unwrap();
790        let current = generation(machine.start().unwrap());
791        let before = format!("{machine:?}");
792        assert_eq!(
793            machine
794                .generated(stale, response(TurnEndReason::ToolUse, 1))
795                .unwrap_err(),
796            StateError::UnexpectedCompletion
797        );
798        assert_eq!(
799            machine.append_input(user("interjection")),
800            Err(StateError::Busy)
801        );
802        assert_eq!(
803            machine.replace_context(vec![user("replacement")], None),
804            Err(StateError::Busy)
805        );
806        assert!(matches!(machine.truncate_history(0), Err(StateError::Busy)));
807        assert_eq!(
808            machine.generation_failed(stale),
809            Err(StateError::UnexpectedCompletion)
810        );
811        assert_eq!(format!("{machine:?}"), before);
812        machine.generation_failed(current).unwrap();
813    }
814
815    #[test]
816    fn tool_batch_requires_matching_count_and_can_only_complete_once() {
817        let mut machine = state();
818        let generation = generation(machine.start().unwrap());
819        let Effect::ExecuteTools { operation, .. } = machine
820            .generated(generation, response(TurnEndReason::ToolUse, 2))
821            .unwrap()
822        else {
823            panic!()
824        };
825        let before = format!("{machine:?}");
826        assert!(
827            machine
828                .tools_completed(operation, vec![ToolResult::text("partial")], false)
829                .is_err()
830        );
831        assert_eq!(format!("{machine:?}"), before);
832        assert!(matches!(
833            machine
834                .tools_completed(
835                    operation,
836                    vec![ToolResult::text("finished"), ToolResult::err("cancelled")],
837                    true
838                )
839                .unwrap(),
840            Effect::Cancelled
841        ));
842        validate_context(machine.history()).unwrap();
843        assert_eq!(
844            machine
845                .tools_completed(operation, vec![], false)
846                .unwrap_err(),
847            StateError::UnexpectedCompletion
848        );
849        machine.append_input(user("next turn")).unwrap();
850    }
851
852    #[test]
853    fn context_replacement_and_rewind_reject_split_or_orphaned_tool_batches() {
854        let mut machine = state();
855        let call = assistant_tool(None, "test", json!({}));
856        for history in [
857            vec![call.clone()],
858            vec![tool_results(&["orphan"])],
859            vec![call.clone(), user("interrupt"), tool_results(&["result"])],
860            vec![call.clone(), tool_results(&["extra", "extra"])],
861            vec![Message::ToolResults {
862                tool_use_results: vec![],
863            }],
864        ] {
865            assert!(machine.replace_context(history, None).is_err());
866            assert_eq!(machine.history().len(), 1);
867        }
868        machine
869            .replace_context(
870                vec![user("task"), call, tool_results(&["ok"]), assistant("done")],
871                None,
872            )
873            .unwrap();
874        assert!(machine.truncate_history(2).is_err());
875        assert!(machine.truncate_history(99).is_err());
876        assert_eq!(machine.history().len(), 4);
877        machine.truncate_history(1).unwrap();
878    }
879
880    #[test]
881    fn truncation_cap_is_an_explicit_finish_reason_and_resets_for_a_new_run() {
882        let mut machine = state();
883        machine.set_max_truncated_resumes(1);
884        let first = generation(machine.start().unwrap());
885        let next = machine
886            .generated(first, response(TurnEndReason::MaxTokens, 0))
887            .unwrap();
888        let second = generation(next);
889        assert!(matches!(
890            machine
891                .generated(second, response(TurnEndReason::MaxTokens, 0))
892                .unwrap(),
893            Effect::Finished {
894                reason: TurnEndReason::MaxTokens,
895                ..
896            }
897        ));
898        machine.append_input(user("new task")).unwrap();
899        let next = generation(machine.start().unwrap());
900        assert!(matches!(
901            machine
902                .generated(next, response(TurnEndReason::MaxTokens, 0))
903                .unwrap(),
904            Effect::Generate { .. }
905        ));
906    }
907
908    #[test]
909    fn restarting_failed_generation_preserves_usage_and_truncation_limits() {
910        let mut machine = state();
911        machine.set_max_truncated_resumes(1);
912        let first = generation(machine.start().unwrap());
913        let pending = generation(
914            machine
915                .generated(first, response(TurnEndReason::MaxTokens, 0))
916                .unwrap(),
917        );
918        machine.generation_failed(pending).unwrap();
919        let retried = generation(machine.start().unwrap());
920        assert_eq!(machine.run_usage().unwrap().output_tokens, 10);
921        assert!(matches!(
922            machine
923                .generated(retried, response(TurnEndReason::MaxTokens, 0))
924                .unwrap(),
925            Effect::Finished {
926                reason: TurnEndReason::MaxTokens,
927                ..
928            }
929        ));
930        assert_eq!(machine.run_usage().unwrap().output_tokens, 20);
931
932        machine.append_input(user("new task")).unwrap();
933        let fresh = generation(machine.start().unwrap());
934        assert!(machine.run_usage().is_none());
935        assert!(matches!(
936            machine
937                .generated(fresh, response(TurnEndReason::MaxTokens, 0))
938                .unwrap(),
939            Effect::Generate { .. }
940        ));
941        assert_eq!(machine.run_usage().unwrap().output_tokens, 10);
942    }
943
944    #[test]
945    fn invalid_outputs_leave_history_usage_and_truncation_budget_unchanged() {
946        let mut machine = state();
947        machine.set_max_truncated_resumes(1);
948        let first = generation(machine.start().unwrap());
949        let mut pending = generation(
950            machine
951                .generated(first, response(TurnEndReason::MaxTokens, 0))
952                .unwrap(),
953        );
954        let history = serde_json::to_value(machine.history()).unwrap();
955        let usage = machine.run_usage();
956        let mut system_output = response(TurnEndReason::EndTurn, 0);
957        system_output.content.push(Content::System {
958            kind: "invalid".into(),
959            text: "invalid".into(),
960            data: serde_json::Value::Null,
961        });
962        for output in [response(TurnEndReason::ToolUse, 0), system_output] {
963            machine.begin_effect(pending).unwrap();
964            assert!(matches!(
965                machine.generated(pending, output),
966                Err(StateError::InvalidResponse(_))
967            ));
968            assert_eq!(serde_json::to_value(machine.history()).unwrap(), history);
969            assert_eq!(machine.run_usage(), usage);
970            assert_eq!(machine.last_usage(), usage);
971            pending = generation(machine.start().unwrap());
972        }
973        assert!(matches!(
974            machine
975                .generated(pending, response(TurnEndReason::MaxTokens, 0))
976                .unwrap(),
977            Effect::Finished {
978                reason: TurnEndReason::MaxTokens,
979                ..
980            }
981        ));
982    }
983}
myco_agent/state.rs

myco_agent/
state.rs
