| 1 | //! Turn context and tracking. |
| 2 | //! |
| 3 | //! A "turn" is one user message and the resulting AI response, |
| 4 | //! including any tool calls that occur. |
| 5 | //! |
| 6 | //! ## Snapshot lifecycle hooks |
| 7 | //! |
| 8 | //! [`pre_turn_snapshot`] and [`post_turn_snapshot`] book-end a turn by |
| 9 | //! taking a workspace-level snapshot into a side git repo (see |
| 10 | //! `crate::snapshot`). They are intentionally non-blocking and |
| 11 | //! non-fatal: any IO error is logged at WARN and swallowed so a busted |
| 12 | //! filesystem or missing `git` binary never derails the agent loop. |
| 13 | //! `/restore N` and the `revert_turn` tool both consume these |
| 14 | //! snapshots. |
| 15 | |
| 16 | use crate::models::Usage; |
| 17 | use crate::snapshot::SnapshotRepo; |
| 18 | use std::path::Path; |
| 19 | use std::time::{Duration, Instant}; |
| 20 | |
| 21 | /// Context for a single turn (user message + AI response). |
| 22 | #[derive(Debug)] |
| 23 | pub struct TurnContext { |
| 24 | /// Turn ID |
| 25 | pub id: String, |
| 26 | |
| 27 | /// When the turn started |
| 28 | #[allow(dead_code)] |
| 29 | pub started_at: Instant, |
| 30 | |
| 31 | /// Current step in the turn (tool call iteration) |
| 32 | pub step: u32, |
| 33 | |
| 34 | /// Maximum steps allowed |
| 35 | pub max_steps: u32, |
| 36 | |
| 37 | /// Tool calls made in this turn |
| 38 | pub tool_calls: Vec<TurnToolCall>, |
| 39 | |
| 40 | /// Whether the turn has been cancelled |
| 41 | #[allow(dead_code)] |
| 42 | pub cancelled: bool, |
| 43 | |
| 44 | /// Usage for this turn |
| 45 | pub usage: Usage, |
| 46 | } |
| 47 | |
| 48 | /// Record of a tool call within a turn. |
| 49 | #[derive(Debug, Clone)] |
| 50 | pub struct TurnToolCall { |
| 51 | pub id: String, |
| 52 | pub name: String, |
| 53 | pub input: serde_json::Value, |
| 54 | pub result: Option<String>, |
| 55 | pub error: Option<String>, |
| 56 | pub duration: Option<Duration>, |
| 57 | } |
| 58 | |
| 59 | impl TurnContext { |
| 60 | /// Create a new turn context |
| 61 | pub fn new(max_steps: u32) -> Self { |
| 62 | Self { |
| 63 | id: uuid::Uuid::new_v4().to_string(), |
| 64 | started_at: Instant::now(), |
| 65 | step: 0, |
| 66 | max_steps, |
| 67 | tool_calls: Vec::new(), |
| 68 | cancelled: false, |
| 69 | usage: Usage { |
| 70 | input_tokens: 0, |
| 71 | output_tokens: 0, |
| 72 | ..Usage::default() |
| 73 | }, |
| 74 | } |
| 75 | } |
| 76 | |
| 77 | /// Increment the step counter |
| 78 | pub fn next_step(&mut self) -> bool { |
| 79 | self.step += 1; |
| 80 | self.step <= self.max_steps |
| 81 | } |
| 82 | |
| 83 | /// Check if the turn has reached max steps |
| 84 | pub fn at_max_steps(&self) -> bool { |
| 85 | self.step >= self.max_steps |
| 86 | } |
| 87 | |
| 88 | /// Record a tool call |
| 89 | pub fn record_tool_call(&mut self, call: TurnToolCall) { |
| 90 | self.tool_calls.push(call); |
| 91 | } |
| 92 | |
| 93 | /// Cancel the turn |
| 94 | #[allow(dead_code)] |
| 95 | pub fn cancel(&mut self) { |
| 96 | self.cancelled = true; |
| 97 | } |
| 98 | |
| 99 | /// Get the elapsed time |
| 100 | #[allow(dead_code)] |
| 101 | pub fn elapsed(&self) -> Duration { |
| 102 | self.started_at.elapsed() |
| 103 | } |
| 104 | |
| 105 | /// Add usage from an API response |
| 106 | pub fn add_usage(&mut self, usage: &Usage) { |
| 107 | self.usage.input_tokens += usage.input_tokens; |
| 108 | self.usage.output_tokens += usage.output_tokens; |
| 109 | self.usage.prompt_cache_hit_tokens = add_optional_usage( |
| 110 | self.usage.prompt_cache_hit_tokens, |
| 111 | usage.prompt_cache_hit_tokens, |
| 112 | ); |
| 113 | self.usage.prompt_cache_miss_tokens = add_optional_usage( |
| 114 | self.usage.prompt_cache_miss_tokens, |
| 115 | usage.prompt_cache_miss_tokens, |
| 116 | ); |
| 117 | self.usage.reasoning_tokens = |
| 118 | add_optional_usage(self.usage.reasoning_tokens, usage.reasoning_tokens); |
| 119 | } |
| 120 | } |
| 121 | |
| 122 | fn add_optional_usage(total: Option<u32>, delta: Option<u32>) -> Option<u32> { |
| 123 | match (total, delta) { |
| 124 | (Some(total), Some(delta)) => Some(total.saturating_add(delta)), |
| 125 | (None, Some(delta)) => Some(delta), |
| 126 | (Some(total), None) => Some(total), |
| 127 | (None, None) => None, |
| 128 | } |
| 129 | } |
| 130 | |
| 131 | /// Take a `pre-turn:<seq>` workspace snapshot. |
| 132 | /// |
| 133 | /// Returns the snapshot SHA on success, `None` on any error. Errors are |
| 134 | /// logged at WARN; the turn loop must not block on this. |
| 135 | pub fn pre_turn_snapshot(workspace: &Path, turn_seq: u64) -> Option<String> { |
| 136 | snapshot_with_label(workspace, &format!("pre-turn:{turn_seq}")) |
| 137 | } |
| 138 | |
| 139 | /// Take a `tool:<call_id>` workspace snapshot, taken before executing a |
| 140 | /// file-modifying tool call (write_file, edit_file, apply_patch). |
| 141 | /// |
| 142 | /// This enables surgical undo: `/undo` can restore to the most recent |
| 143 | /// `tool:<call_id>` snapshot to revert just the last file write. |
| 144 | /// |
| 145 | /// Returns the snapshot SHA on success, `None` on any error. Errors are |
| 146 | /// logged at WARN and are non-fatal. |
| 147 | pub fn pre_tool_snapshot(workspace: &Path, call_id: &str) -> Option<String> { |
| 148 | snapshot_with_label(workspace, &format!("tool:{call_id}")) |
| 149 | } |
| 150 | |
| 151 | /// Take a `post-turn:<seq>` workspace snapshot. Same failure model as |
| 152 | /// [`pre_turn_snapshot`]. |
| 153 | pub fn post_turn_snapshot(workspace: &Path, turn_seq: u64) -> Option<String> { |
| 154 | snapshot_with_label(workspace, &format!("post-turn:{turn_seq}")) |
| 155 | } |
| 156 | |
| 157 | fn snapshot_with_label(workspace: &Path, label: &str) -> Option<String> { |
| 158 | match SnapshotRepo::open_or_init(workspace) { |
| 159 | Ok(repo) => match repo.snapshot(label) { |
| 160 | Ok(id) => Some(id.0), |
| 161 | Err(e) => { |
| 162 | tracing::warn!(target: "snapshot", "snapshot '{label}' failed: {e}"); |
| 163 | None |
| 164 | } |
| 165 | }, |
| 166 | Err(e) => { |
| 167 | tracing::warn!(target: "snapshot", "snapshot repo init failed: {e}"); |
| 168 | None |
| 169 | } |
| 170 | } |
| 171 | } |
| 172 | |
| 173 | impl TurnToolCall { |
| 174 | /// Create a new tool call record |
| 175 | pub fn new(id: String, name: String, input: serde_json::Value) -> Self { |
| 176 | Self { |
| 177 | id, |
| 178 | name, |
| 179 | input, |
| 180 | result: None, |
| 181 | error: None, |
| 182 | duration: None, |
| 183 | } |
| 184 | } |
| 185 | |
| 186 | /// Set the result |
| 187 | pub fn set_result(&mut self, result: String, duration: Duration) { |
| 188 | self.result = Some(result); |
| 189 | self.duration = Some(duration); |
| 190 | } |
| 191 | |
| 192 | /// Set an error |
| 193 | pub fn set_error(&mut self, error: String, duration: Duration) { |
| 194 | self.error = Some(error); |
| 195 | self.duration = Some(duration); |
| 196 | } |
| 197 | } |
| 198 |