| 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::core::events::TurnRoute; |
| 17 | use crate::models::Usage; |
| 18 | use crate::snapshot::SnapshotRepo; |
| 19 | use std::path::Path; |
| 20 | use std::time::{Duration, Instant}; |
| 21 | |
| 22 | /// Context for a single turn (user message + AI response). |
| 23 | #[derive(Debug)] |
| 24 | pub struct TurnContext { |
| 25 | /// Turn ID |
| 26 | pub id: String, |
| 27 | |
| 28 | /// When the turn started |
| 29 | #[allow(dead_code)] |
| 30 | pub started_at: Instant, |
| 31 | |
| 32 | /// Current step in the turn (tool call iteration) |
| 33 | pub step: u32, |
| 34 | |
| 35 | /// Maximum steps allowed |
| 36 | pub max_steps: u32, |
| 37 | |
| 38 | /// Number of tool calls made in this turn. |
| 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 | /// Route facts resolved for this turn but not timestamped until the first |
| 48 | /// provider request is actually dispatched. |
| 49 | pub(crate) pending_route: Option<TurnRoute>, |
| 50 | } |
| 51 | |
| 52 | impl TurnContext { |
| 53 | /// Create a new turn context |
| 54 | pub fn new(max_steps: u32) -> Self { |
| 55 | Self { |
| 56 | id: uuid::Uuid::new_v4().to_string(), |
| 57 | started_at: Instant::now(), |
| 58 | step: 0, |
| 59 | max_steps, |
| 60 | cancelled: false, |
| 61 | usage: Usage { |
| 62 | input_tokens: 0, |
| 63 | output_tokens: 0, |
| 64 | ..Usage::default() |
| 65 | }, |
| 66 | pending_route: None, |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | /// Increment the step counter |
| 71 | pub fn next_step(&mut self) -> bool { |
| 72 | self.step += 1; |
| 73 | self.step <= self.max_steps |
| 74 | } |
| 75 | |
| 76 | /// Check if the turn has reached max steps |
| 77 | pub fn at_max_steps(&self) -> bool { |
| 78 | self.step >= self.max_steps |
| 79 | } |
| 80 | |
| 81 | /// Cancel the turn |
| 82 | #[allow(dead_code)] |
| 83 | pub fn cancel(&mut self) { |
| 84 | self.cancelled = true; |
| 85 | } |
| 86 | |
| 87 | /// Get the elapsed time |
| 88 | #[allow(dead_code)] |
| 89 | pub fn elapsed(&self) -> Duration { |
| 90 | self.started_at.elapsed() |
| 91 | } |
| 92 | |
| 93 | /// Add usage from an API response |
| 94 | pub fn add_usage(&mut self, usage: &Usage) { |
| 95 | self.usage.input_tokens = self.usage.input_tokens.saturating_add(usage.input_tokens); |
| 96 | self.usage.output_tokens = self.usage.output_tokens.saturating_add(usage.output_tokens); |
| 97 | self.usage.prompt_cache_hit_tokens = add_optional_usage( |
| 98 | self.usage.prompt_cache_hit_tokens, |
| 99 | usage.prompt_cache_hit_tokens, |
| 100 | ); |
| 101 | self.usage.prompt_cache_miss_tokens = add_optional_usage( |
| 102 | self.usage.prompt_cache_miss_tokens, |
| 103 | usage.prompt_cache_miss_tokens, |
| 104 | ); |
| 105 | self.usage.prompt_cache_write_tokens = add_optional_usage( |
| 106 | self.usage.prompt_cache_write_tokens, |
| 107 | usage.prompt_cache_write_tokens, |
| 108 | ); |
| 109 | self.usage.reasoning_tokens = |
| 110 | add_optional_usage(self.usage.reasoning_tokens, usage.reasoning_tokens); |
| 111 | self.usage.reasoning_replay_tokens = add_optional_usage( |
| 112 | self.usage.reasoning_replay_tokens, |
| 113 | usage.reasoning_replay_tokens, |
| 114 | ); |
| 115 | if let Some(delta) = usage.server_tool_use.as_ref() { |
| 116 | let total = self.usage.server_tool_use.get_or_insert_default(); |
| 117 | total.code_execution_requests = |
| 118 | add_optional_usage(total.code_execution_requests, delta.code_execution_requests); |
| 119 | total.tool_search_requests = |
| 120 | add_optional_usage(total.tool_search_requests, delta.tool_search_requests); |
| 121 | } |
| 122 | } |
| 123 | } |
| 124 | |
| 125 | fn add_optional_usage(total: Option<u32>, delta: Option<u32>) -> Option<u32> { |
| 126 | match (total, delta) { |
| 127 | (Some(total), Some(delta)) => Some(total.saturating_add(delta)), |
| 128 | (None, Some(delta)) => Some(delta), |
| 129 | (Some(total), None) => Some(total), |
| 130 | (None, None) => None, |
| 131 | } |
| 132 | } |
| 133 | |
| 134 | #[cfg(test)] |
| 135 | mod usage_tests { |
| 136 | use super::*; |
| 137 | use crate::models::ServerToolUsage; |
| 138 | |
| 139 | #[test] |
| 140 | fn add_usage_preserves_replay_and_saturates_server_tool_counters() { |
| 141 | let mut turn = TurnContext::new(2); |
| 142 | turn.add_usage(&Usage { |
| 143 | reasoning_replay_tokens: Some(u32::MAX - 1), |
| 144 | server_tool_use: Some(ServerToolUsage { |
| 145 | code_execution_requests: Some(u32::MAX), |
| 146 | tool_search_requests: Some(2), |
| 147 | }), |
| 148 | ..Usage::default() |
| 149 | }); |
| 150 | turn.add_usage(&Usage { |
| 151 | reasoning_replay_tokens: Some(9), |
| 152 | server_tool_use: Some(ServerToolUsage { |
| 153 | code_execution_requests: Some(1), |
| 154 | tool_search_requests: Some(3), |
| 155 | }), |
| 156 | ..Usage::default() |
| 157 | }); |
| 158 | |
| 159 | assert_eq!(turn.usage.reasoning_replay_tokens, Some(u32::MAX)); |
| 160 | let server = turn.usage.server_tool_use.expect("server tool usage"); |
| 161 | assert_eq!(server.code_execution_requests, Some(u32::MAX)); |
| 162 | assert_eq!(server.tool_search_requests, Some(5)); |
| 163 | } |
| 164 | } |
| 165 | |
| 166 | /// Maximum characters of the user prompt snippet to embed in a snapshot |
| 167 | /// label. Longer prompts are truncated with an ellipsis. |
| 168 | const USER_PROMPT_LABEL_MAX: usize = 100; |
| 169 | |
| 170 | /// Format a snapshot label that includes the user prompt for readability |
| 171 | /// in `/restore` listings. |
| 172 | /// |
| 173 | /// Takes the first line of the prompt (up to `USER_PROMPT_LABEL_MAX` |
| 174 | /// characters) and appends it to the traditional `type:seq` label so |
| 175 | /// users can identify which turn each snapshot belongs to. |
| 176 | pub(crate) fn format_snapshot_label( |
| 177 | prefix: &str, |
| 178 | turn_seq: u64, |
| 179 | user_prompt: Option<&str>, |
| 180 | ) -> String { |
| 181 | let base = format!("{prefix}:{turn_seq}"); |
| 182 | match user_prompt { |
| 183 | None | Some("") => base, |
| 184 | Some(prompt) => match snapshot_label_prompt_snippet(prompt) { |
| 185 | None => base, |
| 186 | Some(snippet) => format!("{base}: {snippet}"), |
| 187 | }, |
| 188 | } |
| 189 | } |
| 190 | |
| 191 | /// The exact prompt snippet [`format_snapshot_label`] embeds after `type:seq`. |
| 192 | /// |
| 193 | /// Read surfaces that want to correlate a recorded prompt back to a restore |
| 194 | /// point must go through this function rather than re-deriving the truncation, |
| 195 | /// so the reader and the writer can never disagree about what a label means. |
| 196 | /// Returns `None` when the prompt contributes no snippet at all. |
| 197 | pub(crate) fn snapshot_label_prompt_snippet(prompt: &str) -> Option<String> { |
| 198 | if prompt.is_empty() { |
| 199 | return None; |
| 200 | } |
| 201 | let first_line = prompt.lines().next().unwrap_or(""); |
| 202 | let truncated: String = first_line.chars().take(USER_PROMPT_LABEL_MAX).collect(); |
| 203 | if truncated.chars().count() < first_line.chars().count() { |
| 204 | Some(format!("{truncated}…")) |
| 205 | } else { |
| 206 | Some(truncated) |
| 207 | } |
| 208 | } |
| 209 | |
| 210 | /// A snapshot label parsed back into its parts. |
| 211 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 212 | pub(crate) struct ParsedSnapshotLabel { |
| 213 | /// `pre-turn`, `post-turn`, `tool`, or whatever prefix produced it. |
| 214 | pub kind: String, |
| 215 | /// The turn sequence for `pre-turn`/`post-turn` labels. `tool` labels |
| 216 | /// carry a call id rather than a sequence, so this stays `None` for them. |
| 217 | pub seq: Option<u64>, |
| 218 | /// The embedded prompt snippet, exactly as |
| 219 | /// [`snapshot_label_prompt_snippet`] produced it. |
| 220 | pub prompt_snippet: Option<String>, |
| 221 | } |
| 222 | |
| 223 | /// Parse a label produced by [`format_snapshot_label`]. |
| 224 | /// |
| 225 | /// This is deliberately total: an unrecognized label still yields a record with |
| 226 | /// the raw text as `kind`, because a read surface must describe what is really |
| 227 | /// stored rather than silently dropping rows it does not recognize. |
| 228 | pub(crate) fn parse_snapshot_label(label: &str) -> ParsedSnapshotLabel { |
| 229 | let (head, snippet) = match label.split_once(": ") { |
| 230 | Some((head, rest)) => (head, Some(rest.to_string())), |
| 231 | None => (label, None), |
| 232 | }; |
| 233 | match head.split_once(':') { |
| 234 | Some((kind, seq)) => ParsedSnapshotLabel { |
| 235 | kind: kind.to_string(), |
| 236 | seq: seq.parse::<u64>().ok(), |
| 237 | prompt_snippet: snippet, |
| 238 | }, |
| 239 | None => ParsedSnapshotLabel { |
| 240 | kind: head.to_string(), |
| 241 | seq: None, |
| 242 | prompt_snippet: snippet, |
| 243 | }, |
| 244 | } |
| 245 | } |
| 246 | |
| 247 | /// Take a `pre-turn:<seq>` workspace snapshot. |
| 248 | /// |
| 249 | /// `cap_bytes` is the workspace-size ceiling that gates first-init |
| 250 | /// (passed through to [`SnapshotRepo::open_or_init_with_cap`]); pass |
| 251 | /// `0` to disable the cap. |
| 252 | /// `user_prompt` is an optional snippet of the user's message for this |
| 253 | /// turn, embedded in the snapshot label so `/restore` listings are |
| 254 | /// human-readable. |
| 255 | /// |
| 256 | /// Returns the snapshot SHA on success, `None` on any error. Errors are |
| 257 | /// logged at WARN; the turn loop must not block on this. |
| 258 | pub fn pre_turn_snapshot( |
| 259 | workspace: &Path, |
| 260 | turn_seq: u64, |
| 261 | cap_bytes: u64, |
| 262 | user_prompt: Option<&str>, |
| 263 | session_id: Option<&str>, |
| 264 | ) -> Option<String> { |
| 265 | snapshot_with_label( |
| 266 | workspace, |
| 267 | &format_snapshot_label("pre-turn", turn_seq, user_prompt), |
| 268 | cap_bytes, |
| 269 | session_id, |
| 270 | ) |
| 271 | } |
| 272 | |
| 273 | /// Take a `tool:<call_id>` workspace snapshot, taken before executing a |
| 274 | /// file-modifying tool call (write_file, edit_file, apply_patch). |
| 275 | /// |
| 276 | /// This enables surgical undo: `/undo` can restore to the most recent |
| 277 | /// `tool:<call_id>` snapshot to revert just the last file write. |
| 278 | /// |
| 279 | /// Returns the snapshot SHA on success, `None` on any error. Errors are |
| 280 | /// logged at WARN and are non-fatal. |
| 281 | pub fn pre_tool_snapshot( |
| 282 | workspace: &Path, |
| 283 | call_id: &str, |
| 284 | cap_bytes: u64, |
| 285 | session_id: Option<&str>, |
| 286 | ) -> Option<String> { |
| 287 | snapshot_with_label(workspace, &format!("tool:{call_id}"), cap_bytes, session_id) |
| 288 | } |
| 289 | |
| 290 | /// Take a `post-turn:<seq>` workspace snapshot. Same failure model as |
| 291 | /// [`pre_turn_snapshot`]. |
| 292 | pub fn post_turn_snapshot( |
| 293 | workspace: &Path, |
| 294 | turn_seq: u64, |
| 295 | cap_bytes: u64, |
| 296 | user_prompt: Option<&str>, |
| 297 | session_id: Option<&str>, |
| 298 | ) -> Option<String> { |
| 299 | snapshot_with_label( |
| 300 | workspace, |
| 301 | &format_snapshot_label("post-turn", turn_seq, user_prompt), |
| 302 | cap_bytes, |
| 303 | session_id, |
| 304 | ) |
| 305 | } |
| 306 | |
| 307 | fn snapshot_with_label( |
| 308 | workspace: &Path, |
| 309 | label: &str, |
| 310 | cap_bytes: u64, |
| 311 | session_id: Option<&str>, |
| 312 | ) -> Option<String> { |
| 313 | match SnapshotRepo::open_or_init_with_cap(workspace, cap_bytes) { |
| 314 | Ok(repo) => { |
| 315 | let id = match repo.snapshot_with_session(label, session_id) { |
| 316 | Ok(id) => Some(id.0), |
| 317 | Err(e) => { |
| 318 | tracing::warn!(target: "snapshot", "snapshot '{label}' failed: {e}"); |
| 319 | return None; |
| 320 | } |
| 321 | }; |
| 322 | // Prune oldest snapshots to cap disk usage (#1112). |
| 323 | if let Err(e) = repo.prune_keep_last_n(crate::snapshot::DEFAULT_MAX_SNAPSHOTS) { |
| 324 | tracing::warn!(target: "snapshot", "snapshot prune failed: {e}"); |
| 325 | } |
| 326 | id |
| 327 | } |
| 328 | Err(e) => { |
| 329 | tracing::warn!(target: "snapshot", "snapshot repo init failed: {e}"); |
| 330 | maybe_notify_snapshots_disabled_once(workspace, &e); |
| 331 | None |
| 332 | } |
| 333 | } |
| 334 | } |
| 335 | |
| 336 | // The stderr print is deliberate: headless/CLI stderr is the user surface for |
| 337 | // this once-per-workspace warning, matching the pre-TUI notices in |
| 338 | // runtime_log.rs. |
| 339 | #[allow(clippy::print_stderr)] |
| 340 | fn maybe_notify_snapshots_disabled_once(workspace: &Path, error: &std::io::Error) { |
| 341 | let message = error.to_string(); |
| 342 | if !(message.contains("workspace too large for snapshots") |
| 343 | || message.contains("workspace snapshots are disabled")) |
| 344 | { |
| 345 | return; |
| 346 | } |
| 347 | use std::collections::HashSet; |
| 348 | use std::sync::{Mutex, OnceLock}; |
| 349 | static NOTIFIED: OnceLock<Mutex<HashSet<String>>> = OnceLock::new(); |
| 350 | let key = workspace.to_string_lossy().into_owned(); |
| 351 | let set = NOTIFIED.get_or_init(|| Mutex::new(HashSet::new())); |
| 352 | let Ok(mut guard) = set.lock() else { |
| 353 | return; |
| 354 | }; |
| 355 | if !guard.insert(key) { |
| 356 | return; |
| 357 | } |
| 358 | // One prominent notice per workspace process lifetime — silent disable is |
| 359 | // the §2.7 failure mode. Opt-in remains `[snapshots] max_workspace_gb` |
| 360 | // (raise the cap or set 0 to disable the size gate). |
| 361 | eprintln!( |
| 362 | "warning: workspace snapshots/undo are OFF for {} |
| 363 | {message} |
| 364 | raise `[snapshots] max_workspace_gb` in config.toml (or set it to 0 to disable the cap) to opt in.", |
| 365 | workspace.display() |
| 366 | ); |
| 367 | } |
| 368 | |
| 369 | #[cfg(test)] |
| 370 | mod snapshot_label_tests { |
| 371 | use super::*; |
| 372 | |
| 373 | #[test] |
| 374 | fn label_writer_and_parser_agree_on_prompt_snippet() { |
| 375 | let prompt = "rename the widget\nsecond line is dropped"; |
| 376 | let label = format_snapshot_label("pre-turn", 7, Some(prompt)); |
| 377 | assert_eq!(label, "pre-turn:7: rename the widget"); |
| 378 | |
| 379 | let parsed = parse_snapshot_label(&label); |
| 380 | assert_eq!(parsed.kind, "pre-turn"); |
| 381 | assert_eq!(parsed.seq, Some(7)); |
| 382 | assert_eq!( |
| 383 | parsed.prompt_snippet.as_deref(), |
| 384 | snapshot_label_prompt_snippet(prompt).as_deref(), |
| 385 | "a reader must recover exactly the snippet the writer embedded" |
| 386 | ); |
| 387 | } |
| 388 | |
| 389 | #[test] |
| 390 | fn truncated_prompt_round_trips_with_its_ellipsis() { |
| 391 | let prompt = "x".repeat(USER_PROMPT_LABEL_MAX + 25); |
| 392 | let label = format_snapshot_label("post-turn", 2, Some(&prompt)); |
| 393 | let parsed = parse_snapshot_label(&label); |
| 394 | let snippet = parsed.prompt_snippet.expect("snippet"); |
| 395 | assert!(snippet.ends_with('…')); |
| 396 | assert_eq!(snippet.chars().count(), USER_PROMPT_LABEL_MAX + 1); |
| 397 | assert_eq!( |
| 398 | Some(snippet), |
| 399 | snapshot_label_prompt_snippet(&prompt), |
| 400 | "truncated snippets must also round-trip" |
| 401 | ); |
| 402 | } |
| 403 | |
| 404 | #[test] |
| 405 | fn labels_without_a_prompt_parse_without_inventing_one() { |
| 406 | let label = format_snapshot_label("pre-turn", 3, None); |
| 407 | assert_eq!(label, "pre-turn:3"); |
| 408 | let parsed = parse_snapshot_label(&label); |
| 409 | assert_eq!(parsed.kind, "pre-turn"); |
| 410 | assert_eq!(parsed.seq, Some(3)); |
| 411 | assert_eq!(parsed.prompt_snippet, None); |
| 412 | } |
| 413 | |
| 414 | #[test] |
| 415 | fn tool_labels_carry_a_call_id_not_a_sequence() { |
| 416 | let label = format!("tool:{}", "call_abc123"); |
| 417 | let parsed = parse_snapshot_label(&label); |
| 418 | assert_eq!(parsed.kind, "tool"); |
| 419 | assert_eq!(parsed.seq, None, "a call id is not a turn sequence"); |
| 420 | assert_eq!(parsed.prompt_snippet, None); |
| 421 | } |
| 422 | |
| 423 | #[test] |
| 424 | fn unrecognized_labels_are_reported_rather_than_dropped() { |
| 425 | let parsed = parse_snapshot_label("manual checkpoint"); |
| 426 | assert_eq!(parsed.kind, "manual checkpoint"); |
| 427 | assert_eq!(parsed.seq, None); |
| 428 | assert_eq!(parsed.prompt_snippet, None); |
| 429 | } |
| 430 | |
| 431 | #[test] |
| 432 | fn empty_prompt_contributes_no_snippet() { |
| 433 | assert_eq!(snapshot_label_prompt_snippet(""), None); |
| 434 | assert_eq!(format_snapshot_label("pre-turn", 1, Some("")), "pre-turn:1"); |
| 435 | } |
| 436 | } |
| 437 |