vtcode_commons/ui_protocol/
progress.rs1use std::sync::atomic::{AtomicU64, Ordering};
2use std::time::Instant;
3
4static NEXT_OPERATION: AtomicU64 = AtomicU64::new(1);
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub struct ProgressOperation {
9 id: u64,
10 started_at: Instant,
11}
12
13impl ProgressOperation {
14 pub fn start() -> Self {
15 Self {
16 id: NEXT_OPERATION.fetch_add(1, Ordering::Relaxed),
17 started_at: Instant::now(),
18 }
19 }
20
21 pub fn id(self) -> u64 {
22 self.id
23 }
24
25 pub fn started_at(self) -> Instant {
26 self.started_at
27 }
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum ProgressPhase {
33 Initializing,
34 PreparingContext,
35 SavingCheckpoint,
36 WaitingForModel,
37 Processing,
38 ReceivingResponse,
39 Retrying,
40 CheckingPermissions,
41 WaitingForApproval,
42 RunningTools,
43 CheckingToolOutput,
44 Compacting {
45 estimated_input_tokens: Option<usize>,
46 started_at: Instant,
47 },
48}
49
50impl ProgressPhase {
51 pub fn label(self) -> &'static str {
52 match self {
53 Self::Initializing => "Initializing session",
54 Self::PreparingContext => "Preparing context",
55 Self::SavingCheckpoint => "Saving checkpoint",
56 Self::WaitingForModel => "Waiting for model",
57 Self::Processing => "Processing",
58 Self::ReceivingResponse => "Receiving response",
59 Self::Retrying => "Retrying request",
60 Self::CheckingPermissions => "Checking permissions",
61 Self::WaitingForApproval => "Waiting for approval",
62 Self::RunningTools => "Running tools",
63 Self::CheckingToolOutput => "Checking tool output...",
64 Self::Compacting { .. } => "Compacting context",
65 }
66 }
67
68 pub fn is_animated(self) -> bool {
69 self != Self::WaitingForApproval
70 }
71
72 pub fn elapsed_secs(self, operation: ProgressOperation, now: Instant) -> u64 {
75 let started_at = match self {
76 Self::Compacting { started_at, .. } => started_at,
77 _ => operation.started_at(),
78 };
79 now.saturating_duration_since(started_at).as_secs()
80 }
81
82 pub fn format(self, elapsed_secs: u64) -> String {
83 if let Self::Compacting { estimated_input_tokens, .. } = self {
84 let estimate = estimated_input_tokens
85 .map(|tokens| format!(" · ~{} input tokens", crate::format_token_count(tokens)))
86 .unwrap_or_default();
87 return if elapsed_secs == 0 {
88 format!("{}{}", self.label(), estimate)
89 } else {
90 format!("{}{} · {elapsed_secs}s", self.label(), estimate)
91 };
92 }
93 if elapsed_secs == 0 {
94 self.label().to_string()
95 } else {
96 format!("{} · {elapsed_secs}s", self.label())
97 }
98 }
99
100 pub fn format_for_width(self, elapsed_secs: u64, width: usize) -> String {
103 let formatted = self.format(elapsed_secs);
104 if formatted.chars().count() <= width {
105 return formatted;
106 }
107
108 let Self::Compacting { estimated_input_tokens, .. } = self else {
109 return formatted;
110 };
111 let compact = match estimated_input_tokens {
112 Some(tokens) => format!("Compact ~{} · {elapsed_secs}s", crate::format_token_count(tokens)),
113 None => format!("Compact · {elapsed_secs}s"),
114 };
115 if compact.chars().count() <= width {
116 return compact;
117 }
118 match estimated_input_tokens {
121 Some(tokens) => format!("Compact ~{} {elapsed_secs}s", crate::format_token_count(tokens)),
122 None => compact,
123 }
124 }
125}
126
127#[derive(Debug, Clone, Copy, PartialEq, Eq)]
128pub enum ProgressUpdate {
129 Begin {
130 operation: ProgressOperation,
131 phase: ProgressPhase,
132 },
133 Phase {
134 operation: ProgressOperation,
135 phase: ProgressPhase,
136 },
137 Finish {
138 operation: ProgressOperation,
139 },
140}
141
142#[cfg(test)]
143mod tests {
144 use super::*;
145 use std::time::Duration;
146
147 #[test]
148 fn compaction_phase_uses_its_own_elapsed_clock() {
149 let operation = ProgressOperation::start();
150 let started_at = operation.started_at() + Duration::from_secs(10);
151 let phase = ProgressPhase::Compacting { estimated_input_tokens: Some(12_000), started_at };
152 let now = started_at + Duration::from_secs(74);
153
154 assert_eq!(phase.elapsed_secs(operation, now), 74);
155 assert_eq!(ProgressPhase::WaitingForModel.elapsed_secs(operation, now), 84);
156 }
157}