ai_agent/bootstrap/
state.rs

1// Source: /data/home/swei/claudecode/openclaudecode/src/bootstrap/state.ts
2#![allow(dead_code)]
3
4use once_cell::sync::Lazy;
5use serde::{Deserialize, Serialize};
6use std::collections::{HashMap, HashSet};
7use std::sync::Mutex;
8use uuid::Uuid;
9
10pub type SessionId = String;
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub enum ChannelEntry {
14    Plugin {
15        name: String,
16        marketplace: String,
17        dev: Option<bool>,
18    },
19    Server {
20        name: String,
21        dev: Option<bool>,
22    },
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct ModelUsage {
27    pub input_tokens: u64,
28    pub output_tokens: u64,
29    pub cache_read_input_tokens: u64,
30    pub cache_creation_input_tokens: u64,
31    pub web_search_requests: u64,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct ModelSetting {
36    pub value: Option<String>,
37    pub label: String,
38    pub description: String,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct ModelStrings {
43    pub region_string: Option<String>,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ErrorLogEntry {
48    pub error: String,
49    pub timestamp: String,
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct SessionCronTask {
54    pub id: String,
55    pub cron: String,
56    pub prompt: String,
57    pub created_at: u64,
58    pub recurring: Option<bool>,
59    pub agent_id: Option<String>,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct TeleportedSessionInfo {
64    pub is_teleported: bool,
65    pub has_logged_first_message: bool,
66    pub session_id: Option<String>,
67}
68
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct InvokedSkillInfo {
71    pub skill_name: String,
72    pub skill_path: String,
73    pub content: String,
74    pub invoked_at: u64,
75    pub agent_id: Option<String>,
76}
77
78#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct SlowOperation {
80    pub operation: String,
81    pub duration_ms: f64,
82    pub timestamp: u64,
83}
84
85#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct HookMatcher {
87    pub name: Option<String>,
88    pub matcher: Option<serde_json::Value>,
89    pub plugin_root: Option<String>,
90}
91
92pub struct CostStateRestoreParams {
93    pub total_cost_usd: f64,
94    pub total_api_duration: f64,
95    pub total_api_duration_without_retries: f64,
96    pub total_tool_duration: f64,
97    pub total_lines_added: u64,
98    pub total_lines_removed: u64,
99    pub last_duration: Option<u64>,
100    pub model_usage: Option<HashMap<String, ModelUsage>>,
101}
102
103#[derive(Default)]
104pub struct RegenerateOptions {
105    pub set_current_as_parent: bool,
106}
107
108fn current_timestamp() -> u64 {
109    std::time::SystemTime::now()
110        .duration_since(std::time::UNIX_EPOCH)
111        .map(|d| d.as_millis() as u64)
112        .unwrap_or(0)
113}
114
115pub fn get_session_id() -> String {
116    STATE.lock().unwrap().session_id.clone()
117}
118
119pub fn regenerate_session_id(options: RegenerateOptions) -> String {
120    let mut state = STATE.lock().unwrap();
121    if options.set_current_as_parent {
122        state.parent_session_id = Some(state.session_id.clone());
123    }
124    let old_session_id = state.session_id.clone();
125    state.plan_slug_cache.remove(&old_session_id);
126    let new_id = Uuid::new_v4().to_string();
127    state.session_id = new_id.clone();
128    state.session_project_dir = None;
129    new_id
130}
131
132pub fn get_parent_session_id() -> Option<String> {
133    STATE.lock().unwrap().parent_session_id.clone()
134}
135
136pub fn switch_session(session_id: String, project_dir: Option<String>) {
137    let mut state = STATE.lock().unwrap();
138    let old_session_id = state.session_id.clone();
139    state.plan_slug_cache.remove(&old_session_id);
140    state.session_id = session_id;
141    state.session_project_dir = project_dir;
142}
143
144pub fn get_session_project_dir() -> Option<String> {
145    STATE.lock().unwrap().session_project_dir.clone()
146}
147
148pub fn get_original_cwd() -> String {
149    STATE.lock().unwrap().original_cwd.clone()
150}
151
152pub fn get_project_root() -> String {
153    STATE.lock().unwrap().project_root.clone()
154}
155
156pub fn set_original_cwd(cwd: String) {
157    STATE.lock().unwrap().original_cwd = cwd;
158}
159
160pub fn set_project_root(cwd: String) {
161    STATE.lock().unwrap().project_root = cwd;
162}
163
164pub fn get_cwd_state() -> String {
165    STATE.lock().unwrap().cwd.clone()
166}
167
168pub fn set_cwd_state(cwd: String) {
169    STATE.lock().unwrap().cwd = cwd;
170}
171
172pub fn get_direct_connect_server_url() -> Option<String> {
173    STATE.lock().unwrap().direct_connect_server_url.clone()
174}
175
176pub fn set_direct_connect_server_url(url: String) {
177    STATE.lock().unwrap().direct_connect_server_url = Some(url);
178}
179
180pub fn add_to_total_duration_state(duration: f64, duration_without_retries: f64) {
181    let mut state = STATE.lock().unwrap();
182    state.total_api_duration += duration;
183    state.total_api_duration_without_retries += duration_without_retries;
184}
185
186pub fn reset_total_duration_state_and_cost_for_tests_only() {
187    let mut state = STATE.lock().unwrap();
188    state.total_api_duration = 0.0;
189    state.total_api_duration_without_retries = 0.0;
190    state.total_cost_usd = 0.0;
191}
192
193pub fn add_to_total_cost_state(cost: f64, model_usage: ModelUsage, model: String) {
194    let mut state = STATE.lock().unwrap();
195    state.model_usage.insert(model, model_usage);
196    state.total_cost_usd += cost;
197}
198
199pub fn get_total_cost_usd() -> f64 {
200    STATE.lock().unwrap().total_cost_usd
201}
202
203pub fn get_total_api_duration() -> f64 {
204    STATE.lock().unwrap().total_api_duration
205}
206
207pub fn get_total_duration() -> u64 {
208    current_timestamp() - STATE.lock().unwrap().start_time
209}
210
211pub fn get_total_api_duration_without_retries() -> f64 {
212    STATE.lock().unwrap().total_api_duration_without_retries
213}
214
215pub fn get_total_tool_duration() -> f64 {
216    STATE.lock().unwrap().total_tool_duration
217}
218
219pub fn add_to_tool_duration(duration: f64) {
220    let mut state = STATE.lock().unwrap();
221    state.total_tool_duration += duration;
222    state.turn_tool_duration_ms += duration;
223    state.turn_tool_count += 1;
224}
225
226pub fn get_turn_hook_duration_ms() -> f64 {
227    STATE.lock().unwrap().turn_hook_duration_ms
228}
229
230pub fn add_to_turn_hook_duration(duration: f64) {
231    let mut state = STATE.lock().unwrap();
232    state.turn_hook_duration_ms += duration;
233    state.turn_hook_count += 1;
234}
235
236pub fn reset_turn_hook_duration() {
237    let mut state = STATE.lock().unwrap();
238    state.turn_hook_duration_ms = 0.0;
239    state.turn_hook_count = 0;
240}
241
242pub fn get_turn_hook_count() -> u64 {
243    STATE.lock().unwrap().turn_hook_count
244}
245
246pub fn get_turn_tool_duration_ms() -> f64 {
247    STATE.lock().unwrap().turn_tool_duration_ms
248}
249
250pub fn reset_turn_tool_duration() {
251    let mut state = STATE.lock().unwrap();
252    state.turn_tool_duration_ms = 0.0;
253    state.turn_tool_count = 0;
254}
255
256pub fn get_turn_tool_count() -> u64 {
257    STATE.lock().unwrap().turn_tool_count
258}
259
260pub fn get_turn_classifier_duration_ms() -> f64 {
261    STATE.lock().unwrap().turn_classifier_duration_ms
262}
263
264pub fn add_to_turn_classifier_duration(duration: f64) {
265    let mut state = STATE.lock().unwrap();
266    state.turn_classifier_duration_ms += duration;
267    state.turn_classifier_count += 1;
268}
269
270pub fn reset_turn_classifier_duration() {
271    let mut state = STATE.lock().unwrap();
272    state.turn_classifier_duration_ms = 0.0;
273    state.turn_classifier_count = 0;
274}
275
276pub fn get_turn_classifier_count() -> u64 {
277    STATE.lock().unwrap().turn_classifier_count
278}
279
280pub fn get_last_interaction_time() -> u64 {
281    STATE.lock().unwrap().last_interaction_time
282}
283
284pub fn add_to_total_lines_changed(added: u64, removed: u64) {
285    let mut state = STATE.lock().unwrap();
286    state.total_lines_added += added;
287    state.total_lines_removed += removed;
288}
289
290pub fn get_total_lines_added() -> u64 {
291    STATE.lock().unwrap().total_lines_added
292}
293
294pub fn get_total_lines_removed() -> u64 {
295    STATE.lock().unwrap().total_lines_removed
296}
297
298pub fn get_total_input_tokens() -> u64 {
299    STATE
300        .lock()
301        .unwrap()
302        .model_usage
303        .values()
304        .map(|u| u.input_tokens)
305        .sum()
306}
307
308pub fn get_total_output_tokens() -> u64 {
309    STATE
310        .lock()
311        .unwrap()
312        .model_usage
313        .values()
314        .map(|u| u.output_tokens)
315        .sum()
316}
317
318pub fn get_total_cache_read_input_tokens() -> u64 {
319    STATE
320        .lock()
321        .unwrap()
322        .model_usage
323        .values()
324        .map(|u| u.cache_read_input_tokens)
325        .sum()
326}
327
328pub fn get_total_cache_creation_input_tokens() -> u64 {
329    STATE
330        .lock()
331        .unwrap()
332        .model_usage
333        .values()
334        .map(|u| u.cache_creation_input_tokens)
335        .sum()
336}
337
338pub fn get_total_web_search_requests() -> u64 {
339    STATE
340        .lock()
341        .unwrap()
342        .model_usage
343        .values()
344        .map(|u| u.web_search_requests)
345        .sum()
346}
347
348pub fn get_turn_output_tokens() -> u64 {
349    let state = STATE.lock().unwrap();
350    let total = state.model_usage.values().map(|u| u.output_tokens).sum::<u64>();
351    total.saturating_sub(state.output_tokens_at_turn_start)
352}
353
354/// Per-turn output token snapshotting (matches TypeScript bootstrap/state.ts)
355pub fn snapshot_output_tokens_for_turn(budget: Option<f64>) {
356    let mut state = STATE.lock().unwrap();
357    state.output_tokens_at_turn_start = state.model_usage.values().map(|u| u.output_tokens).sum();
358    state.current_turn_token_budget = budget;
359    state.budget_continuation_count = 0;
360}
361
362pub fn get_current_turn_token_budget() -> Option<f64> {
363    STATE.lock().unwrap().current_turn_token_budget
364}
365
366pub fn get_budget_continuation_count() -> u64 {
367    STATE.lock().unwrap().budget_continuation_count
368}
369
370pub fn increment_budget_continuation_count() {
371    let mut state = STATE.lock().unwrap();
372    state.budget_continuation_count += 1;
373}
374
375pub fn set_has_unknown_model_cost() {
376    STATE.lock().unwrap().has_unknown_model_cost = true;
377}
378
379pub fn has_unknown_model_cost() -> bool {
380    STATE.lock().unwrap().has_unknown_model_cost
381}
382
383pub fn get_last_main_request_id() -> Option<String> {
384    STATE.lock().unwrap().last_main_request_id.clone()
385}
386
387pub fn set_last_main_request_id(request_id: String) {
388    STATE.lock().unwrap().last_main_request_id = Some(request_id);
389}
390
391pub fn get_last_api_completion_timestamp() -> Option<u64> {
392    STATE.lock().unwrap().last_api_completion_timestamp
393}
394
395pub fn set_last_api_completion_timestamp(timestamp: u64) {
396    STATE.lock().unwrap().last_api_completion_timestamp = Some(timestamp);
397}
398
399pub fn mark_post_compaction() {
400    STATE.lock().unwrap().pending_post_compaction = true;
401}
402
403pub fn consume_post_compaction() -> bool {
404    let mut state = STATE.lock().unwrap();
405    let was = state.pending_post_compaction;
406    state.pending_post_compaction = false;
407    was
408}
409
410pub fn get_model_usage() -> HashMap<String, ModelUsage> {
411    STATE.lock().unwrap().model_usage.clone()
412}
413
414pub fn get_usage_for_model(model: &str) -> Option<ModelUsage> {
415    STATE.lock().unwrap().model_usage.get(model).cloned()
416}
417
418pub fn get_main_loop_model_override() -> Option<ModelSetting> {
419    STATE.lock().unwrap().main_loop_model_override.clone()
420}
421
422pub fn get_initial_main_loop_model() -> Option<ModelSetting> {
423    STATE.lock().unwrap().initial_main_loop_model.clone()
424}
425
426pub fn set_main_loop_model_override(model: Option<ModelSetting>) {
427    STATE.lock().unwrap().main_loop_model_override = model;
428}
429
430pub fn set_initial_main_loop_model(model: ModelSetting) {
431    STATE.lock().unwrap().initial_main_loop_model = Some(model);
432}
433
434pub fn get_sdk_betas() -> Option<Vec<String>> {
435    STATE.lock().unwrap().sdk_betas.clone()
436}
437
438pub fn set_sdk_betas(betas: Option<Vec<String>>) {
439    STATE.lock().unwrap().sdk_betas = betas;
440}
441
442pub fn reset_cost_state() {
443    let mut state = STATE.lock().unwrap();
444    state.total_cost_usd = 0.0;
445    state.total_api_duration = 0.0;
446    state.total_api_duration_without_retries = 0.0;
447    state.total_tool_duration = 0.0;
448    state.start_time = current_timestamp();
449    state.total_lines_added = 0;
450    state.total_lines_removed = 0;
451    state.has_unknown_model_cost = false;
452    state.model_usage.clear();
453    state.prompt_id = None;
454}
455
456pub fn set_cost_state_for_restore(params: CostStateRestoreParams) {
457    let mut state = STATE.lock().unwrap();
458    state.total_cost_usd = params.total_cost_usd;
459    state.total_api_duration = params.total_api_duration;
460    state.total_api_duration_without_retries = params.total_api_duration_without_retries;
461    state.total_tool_duration = params.total_tool_duration;
462    state.total_lines_added = params.total_lines_added;
463    state.total_lines_removed = params.total_lines_removed;
464
465    if let Some(model_usage) = params.model_usage {
466        state.model_usage = model_usage;
467    }
468
469    if let Some(last_duration) = params.last_duration {
470        state.start_time = current_timestamp() - last_duration;
471    }
472}
473
474pub fn get_model_strings() -> Option<ModelStrings> {
475    STATE.lock().unwrap().model_strings.clone()
476}
477
478pub fn set_model_strings(model_strings: ModelStrings) {
479    STATE.lock().unwrap().model_strings = Some(model_strings);
480}
481
482pub fn reset_model_strings_for_testing_only() {
483    STATE.lock().unwrap().model_strings = None;
484}
485
486pub fn get_is_non_interactive_session() -> bool {
487    !STATE.lock().unwrap().is_interactive
488}
489
490pub fn get_is_interactive() -> bool {
491    STATE.lock().unwrap().is_interactive
492}
493
494pub fn set_is_interactive(value: bool) {
495    STATE.lock().unwrap().is_interactive = value;
496}
497
498pub fn get_client_type() -> String {
499    STATE.lock().unwrap().client_type.clone()
500}
501
502pub fn set_client_type(type_: String) {
503    STATE.lock().unwrap().client_type = type_;
504}
505
506pub fn get_sdk_agent_progress_summaries_enabled() -> bool {
507    STATE.lock().unwrap().sdk_agent_progress_summaries_enabled
508}
509
510pub fn set_sdk_agent_progress_summaries_enabled(value: bool) {
511    STATE.lock().unwrap().sdk_agent_progress_summaries_enabled = value;
512}
513
514pub fn get_kairos_active() -> bool {
515    STATE.lock().unwrap().kairos_active
516}
517
518pub fn set_kairos_active(value: bool) {
519    STATE.lock().unwrap().kairos_active = value;
520}
521
522pub fn get_strict_tool_result_pairing() -> bool {
523    STATE.lock().unwrap().strict_tool_result_pairing
524}
525
526pub fn set_strict_tool_result_pairing(value: bool) {
527    STATE.lock().unwrap().strict_tool_result_pairing = value;
528}
529
530pub fn get_user_msg_opt_in() -> bool {
531    STATE.lock().unwrap().user_msg_opt_in
532}
533
534pub fn set_user_msg_opt_in(value: bool) {
535    STATE.lock().unwrap().user_msg_opt_in = value;
536}
537
538pub fn get_session_source() -> Option<String> {
539    STATE.lock().unwrap().session_source.clone()
540}
541
542pub fn set_session_source(source: String) {
543    STATE.lock().unwrap().session_source = Some(source);
544}
545
546pub fn get_question_preview_format() -> Option<String> {
547    STATE.lock().unwrap().question_preview_format.clone()
548}
549
550pub fn set_question_preview_format(format: String) {
551    STATE.lock().unwrap().question_preview_format = Some(format);
552}
553
554pub fn get_agent_color_map() -> HashMap<String, String> {
555    STATE.lock().unwrap().agent_color_map.clone()
556}
557
558pub fn get_flag_settings_path() -> Option<String> {
559    STATE.lock().unwrap().flag_settings_path.clone()
560}
561
562pub fn set_flag_settings_path(path: Option<String>) {
563    STATE.lock().unwrap().flag_settings_path = path;
564}
565
566pub fn get_flag_settings_inline() -> Option<HashMap<String, serde_json::Value>> {
567    STATE.lock().unwrap().flag_settings_inline.clone()
568}
569
570pub fn set_flag_settings_inline(settings: Option<HashMap<String, serde_json::Value>>) {
571    STATE.lock().unwrap().flag_settings_inline = settings;
572}
573
574pub fn get_session_ingress_token() -> Option<String> {
575    STATE.lock().unwrap().session_ingress_token.clone()
576}
577
578pub fn set_session_ingress_token(token: Option<String>) {
579    STATE.lock().unwrap().session_ingress_token = token;
580}
581
582pub fn get_oauth_token_from_fd() -> Option<String> {
583    STATE.lock().unwrap().oauth_token_from_fd.clone()
584}
585
586pub fn set_oauth_token_from_fd(token: Option<String>) {
587    STATE.lock().unwrap().oauth_token_from_fd = token;
588}
589
590pub fn get_api_key_from_fd() -> Option<String> {
591    STATE.lock().unwrap().api_key_from_fd.clone()
592}
593
594pub fn set_api_key_from_fd(key: Option<String>) {
595    STATE.lock().unwrap().api_key_from_fd = key;
596}
597
598pub fn set_last_api_request(params: Option<serde_json::Value>) {
599    STATE.lock().unwrap().last_api_request = params;
600}
601
602pub fn get_last_api_request() -> Option<serde_json::Value> {
603    STATE.lock().unwrap().last_api_request.clone()
604}
605
606pub fn set_last_api_request_messages(messages: Option<serde_json::Value>) {
607    STATE.lock().unwrap().last_api_request_messages = messages;
608}
609
610pub fn get_last_api_request_messages() -> Option<serde_json::Value> {
611    STATE.lock().unwrap().last_api_request_messages.clone()
612}
613
614pub fn set_last_classifier_requests(requests: Option<Vec<serde_json::Value>>) {
615    STATE.lock().unwrap().last_classifier_requests = requests;
616}
617
618pub fn get_last_classifier_requests() -> Option<Vec<serde_json::Value>> {
619    STATE.lock().unwrap().last_classifier_requests.clone()
620}
621
622pub fn set_cached_claude_md_content(content: Option<String>) {
623    STATE.lock().unwrap().cached_claude_md_content = content;
624}
625
626pub fn get_cached_claude_md_content() -> Option<String> {
627    STATE.lock().unwrap().cached_claude_md_content.clone()
628}
629
630pub fn add_to_in_memory_error_log(error_info: ErrorLogEntry) {
631    const MAX_IN_MEMORY_ERRORS: usize = 100;
632    let mut state = STATE.lock().unwrap();
633    if state.in_memory_error_log.len() >= MAX_IN_MEMORY_ERRORS {
634        state.in_memory_error_log.remove(0);
635    }
636    state.in_memory_error_log.push(error_info);
637}
638
639pub fn get_allowed_setting_sources() -> Vec<String> {
640    STATE.lock().unwrap().allowed_setting_sources.clone()
641}
642
643pub fn set_allowed_setting_sources(sources: Vec<String>) {
644    STATE.lock().unwrap().allowed_setting_sources = sources;
645}
646
647pub fn prefer_third_party_authentication() -> bool {
648    let state = STATE.lock().unwrap();
649    !state.is_interactive && state.client_type != "claude-vscode"
650}
651
652pub fn set_inline_plugins(plugins: Vec<String>) {
653    STATE.lock().unwrap().inline_plugins = plugins;
654}
655
656pub fn get_inline_plugins() -> Vec<String> {
657    STATE.lock().unwrap().inline_plugins.clone()
658}
659
660pub fn set_chrome_flag_override(value: Option<bool>) {
661    STATE.lock().unwrap().chrome_flag_override = value;
662}
663
664pub fn get_chrome_flag_override() -> Option<bool> {
665    STATE.lock().unwrap().chrome_flag_override
666}
667
668pub fn set_use_cowork_plugins(value: bool) {
669    STATE.lock().unwrap().use_cowork_plugins = value;
670}
671
672pub fn get_use_cowork_plugins() -> bool {
673    STATE.lock().unwrap().use_cowork_plugins
674}
675
676pub fn set_session_bypass_permissions_mode(enabled: bool) {
677    STATE.lock().unwrap().session_bypass_permissions_mode = enabled;
678}
679
680pub fn get_session_bypass_permissions_mode() -> bool {
681    STATE.lock().unwrap().session_bypass_permissions_mode
682}
683
684pub fn set_scheduled_tasks_enabled(enabled: bool) {
685    STATE.lock().unwrap().scheduled_tasks_enabled = enabled;
686}
687
688pub fn get_scheduled_tasks_enabled() -> bool {
689    STATE.lock().unwrap().scheduled_tasks_enabled
690}
691
692pub fn get_session_cron_tasks() -> Vec<SessionCronTask> {
693    STATE.lock().unwrap().session_cron_tasks.clone()
694}
695
696pub fn add_session_cron_task(task: SessionCronTask) {
697    STATE.lock().unwrap().session_cron_tasks.push(task);
698}
699
700pub fn remove_session_cron_tasks(ids: &[String]) -> usize {
701    if ids.is_empty() {
702        return 0;
703    }
704    let mut state = STATE.lock().unwrap();
705    let id_set: HashSet<String> = ids.iter().cloned().collect();
706    let initial_len = state.session_cron_tasks.len();
707    state.session_cron_tasks.retain(|t| !id_set.contains(&t.id));
708    let removed = initial_len - state.session_cron_tasks.len();
709    if removed == 0 {
710        return 0;
711    }
712    removed
713}
714
715pub fn set_session_trust_accepted(accepted: bool) {
716    STATE.lock().unwrap().session_trust_accepted = accepted;
717}
718
719pub fn get_session_trust_accepted() -> bool {
720    STATE.lock().unwrap().session_trust_accepted
721}
722
723pub fn set_session_persistence_disabled(disabled: bool) {
724    STATE.lock().unwrap().session_persistence_disabled = disabled;
725}
726
727pub fn is_session_persistence_disabled() -> bool {
728    STATE.lock().unwrap().session_persistence_disabled
729}
730
731pub fn has_exited_plan_mode_in_session() -> bool {
732    STATE.lock().unwrap().has_exited_plan_mode
733}
734
735pub fn set_has_exited_plan_mode(value: bool) {
736    STATE.lock().unwrap().has_exited_plan_mode = value;
737}
738
739pub fn needs_plan_mode_exit_attachment() -> bool {
740    STATE.lock().unwrap().needs_plan_mode_exit_attachment
741}
742
743pub fn set_needs_plan_mode_exit_attachment(value: bool) {
744    STATE.lock().unwrap().needs_plan_mode_exit_attachment = value;
745}
746
747pub fn handle_plan_mode_transition(from_mode: &str, to_mode: &str) {
748    let mut state = STATE.lock().unwrap();
749    if to_mode == "plan" && from_mode != "plan" {
750        state.needs_plan_mode_exit_attachment = false;
751    }
752    if from_mode == "plan" && to_mode != "plan" {
753        state.needs_plan_mode_exit_attachment = true;
754    }
755}
756
757pub fn needs_auto_mode_exit_attachment() -> bool {
758    STATE.lock().unwrap().needs_auto_mode_exit_attachment
759}
760
761pub fn set_needs_auto_mode_exit_attachment(value: bool) {
762    STATE.lock().unwrap().needs_auto_mode_exit_attachment = value;
763}
764
765pub fn handle_auto_mode_transition(from_mode: &str, to_mode: &str) {
766    let mut state = STATE.lock().unwrap();
767    if (from_mode == "auto" && to_mode == "plan") || (from_mode == "plan" && to_mode == "auto") {
768        return;
769    }
770    let from_is_auto = from_mode == "auto";
771    let to_is_auto = to_mode == "auto";
772
773    if to_is_auto && !from_is_auto {
774        state.needs_auto_mode_exit_attachment = false;
775    }
776    if from_is_auto && !to_is_auto {
777        state.needs_auto_mode_exit_attachment = true;
778    }
779}
780
781pub fn has_shown_lsp_recommendation_this_session() -> bool {
782    STATE.lock().unwrap().lsp_recommendation_shown_this_session
783}
784
785pub fn set_lsp_recommendation_shown_this_session(value: bool) {
786    STATE.lock().unwrap().lsp_recommendation_shown_this_session = value;
787}
788
789pub fn set_init_json_schema(schema: HashMap<String, serde_json::Value>) {
790    STATE.lock().unwrap().init_json_schema = Some(schema);
791}
792
793pub fn get_init_json_schema() -> Option<HashMap<String, serde_json::Value>> {
794    STATE.lock().unwrap().init_json_schema.clone()
795}
796
797pub fn get_plan_slug_cache() -> HashMap<String, String> {
798    STATE.lock().unwrap().plan_slug_cache.clone()
799}
800
801pub fn get_session_created_teams() -> HashSet<String> {
802    STATE.lock().unwrap().session_created_teams.clone()
803}
804
805pub fn set_teleported_session_info(info: TeleportedSessionInfo) {
806    STATE.lock().unwrap().teleported_session_info = Some(info);
807}
808
809pub fn get_teleported_session_info() -> Option<TeleportedSessionInfo> {
810    STATE.lock().unwrap().teleported_session_info.clone()
811}
812
813pub fn mark_first_teleport_message_logged() {
814    let mut state = STATE.lock().unwrap();
815    if let Some(info) = state.teleported_session_info.as_mut() {
816        info.has_logged_first_message = true;
817    }
818}
819
820pub fn add_invoked_skill(
821    skill_name: String,
822    skill_path: String,
823    content: String,
824    agent_id: Option<String>,
825) {
826    let key = format!(
827        "{}:{}",
828        agent_id.as_ref().unwrap_or(&String::new()),
829        skill_name
830    );
831    let mut state = STATE.lock().unwrap();
832    state.invoked_skills.insert(
833        key,
834        InvokedSkillInfo {
835            skill_name,
836            skill_path,
837            content,
838            invoked_at: current_timestamp(),
839            agent_id,
840        },
841    );
842}
843
844pub fn get_invoked_skills() -> HashMap<String, InvokedSkillInfo> {
845    STATE.lock().unwrap().invoked_skills.clone()
846}
847
848pub fn get_invoked_skills_for_agent(agent_id: Option<&str>) -> HashMap<String, InvokedSkillInfo> {
849    let normalized_id = agent_id.map(|s| s.to_string());
850    STATE
851        .lock()
852        .unwrap()
853        .invoked_skills
854        .iter()
855        .filter(|(_, skill)| skill.agent_id == normalized_id)
856        .map(|(k, v)| (k.clone(), v.clone()))
857        .collect()
858}
859
860pub fn clear_invoked_skills(preserved_agent_ids: Option<&HashSet<String>>) {
861    let mut state = STATE.lock().unwrap();
862    if let Some(ids) = preserved_agent_ids {
863        if ids.is_empty() {
864            state.invoked_skills.clear();
865            return;
866        }
867        state.invoked_skills.retain(|_, skill| {
868            skill.agent_id.is_none() || !ids.contains(skill.agent_id.as_ref().unwrap())
869        });
870    } else {
871        state.invoked_skills.clear();
872    }
873}
874
875pub fn clear_invoked_skills_for_agent(agent_id: &str) {
876    let mut state = STATE.lock().unwrap();
877    state
878        .invoked_skills
879        .retain(|_, skill| skill.agent_id.as_deref() != Some(agent_id));
880}
881
882const MAX_SLOW_OPERATIONS: usize = 10;
883const SLOW_OPERATION_TTL_MS: u64 = 10000;
884
885pub fn add_slow_operation(operation: String, duration_ms: f64) {
886    let mut state = STATE.lock().unwrap();
887    let now = current_timestamp();
888    state
889        .slow_operations
890        .retain(|op| now - op.timestamp < SLOW_OPERATION_TTL_MS);
891    state.slow_operations.push(SlowOperation {
892        operation,
893        duration_ms,
894        timestamp: now,
895    });
896    if state.slow_operations.len() > MAX_SLOW_OPERATIONS {
897        let len = state.slow_operations.len();
898        state.slow_operations = state.slow_operations.split_off(len - MAX_SLOW_OPERATIONS);
899    }
900}
901
902pub fn get_slow_operations() -> Vec<SlowOperation> {
903    let state = STATE.lock().unwrap();
904    if state.slow_operations.is_empty() {
905        return vec![];
906    }
907    let now = current_timestamp();
908    if state
909        .slow_operations
910        .iter()
911        .any(|op| now - op.timestamp >= SLOW_OPERATION_TTL_MS)
912    {
913        return vec![];
914    }
915    state.slow_operations.clone()
916}
917
918pub fn get_main_thread_agent_type() -> Option<String> {
919    STATE.lock().unwrap().main_thread_agent_type.clone()
920}
921
922pub fn set_main_thread_agent_type(agent_type: Option<String>) {
923    STATE.lock().unwrap().main_thread_agent_type = agent_type;
924}
925
926pub fn get_is_remote_mode() -> bool {
927    STATE.lock().unwrap().is_remote_mode
928}
929
930pub fn set_is_remote_mode(value: bool) {
931    STATE.lock().unwrap().is_remote_mode = value;
932}
933
934pub fn get_system_prompt_section_cache() -> HashMap<String, Option<String>> {
935    STATE.lock().unwrap().system_prompt_section_cache.clone()
936}
937
938pub fn set_system_prompt_section_cache_entry(name: String, value: Option<String>) {
939    STATE
940        .lock()
941        .unwrap()
942        .system_prompt_section_cache
943        .insert(name, value);
944}
945
946pub fn clear_system_prompt_section_state() {
947    STATE.lock().unwrap().system_prompt_section_cache.clear();
948}
949
950pub fn get_last_emitted_date() -> Option<String> {
951    STATE.lock().unwrap().last_emitted_date.clone()
952}
953
954pub fn set_last_emitted_date(date: Option<String>) {
955    STATE.lock().unwrap().last_emitted_date = date;
956}
957
958pub fn get_additional_directories_for_claude_md() -> Vec<String> {
959    STATE
960        .lock()
961        .unwrap()
962        .additional_directories_for_claude_md
963        .clone()
964}
965
966pub fn set_additional_directories_for_claude_md(directories: Vec<String>) {
967    STATE.lock().unwrap().additional_directories_for_claude_md = directories;
968}
969
970pub fn get_allowed_channels() -> Vec<ChannelEntry> {
971    STATE.lock().unwrap().allowed_channels.clone()
972}
973
974pub fn set_allowed_channels(entries: Vec<ChannelEntry>) {
975    STATE.lock().unwrap().allowed_channels = entries;
976}
977
978pub fn get_has_dev_channels() -> bool {
979    STATE.lock().unwrap().has_dev_channels
980}
981
982pub fn set_has_dev_channels(value: bool) {
983    STATE.lock().unwrap().has_dev_channels = value;
984}
985
986pub fn get_prompt_cache_1h_allowlist() -> Option<Vec<String>> {
987    STATE.lock().unwrap().prompt_cache_1h_allowlist.clone()
988}
989
990pub fn set_prompt_cache_1h_allowlist(allowlist: Option<Vec<String>>) {
991    STATE.lock().unwrap().prompt_cache_1h_allowlist = allowlist;
992}
993
994pub fn get_prompt_cache_1h_eligible() -> Option<bool> {
995    STATE.lock().unwrap().prompt_cache_1h_eligible
996}
997
998pub fn set_prompt_cache_1h_eligible(eligible: Option<bool>) {
999    STATE.lock().unwrap().prompt_cache_1h_eligible = eligible;
1000}
1001
1002pub fn get_afk_mode_header_latched() -> Option<bool> {
1003    STATE.lock().unwrap().afk_mode_header_latched
1004}
1005
1006pub fn set_afk_mode_header_latched(v: bool) {
1007    STATE.lock().unwrap().afk_mode_header_latched = Some(v);
1008}
1009
1010pub fn get_fast_mode_header_latched() -> Option<bool> {
1011    STATE.lock().unwrap().fast_mode_header_latched
1012}
1013
1014pub fn set_fast_mode_header_latched(v: bool) {
1015    STATE.lock().unwrap().fast_mode_header_latched = Some(v);
1016}
1017
1018pub fn get_cache_editing_header_latched() -> Option<bool> {
1019    STATE.lock().unwrap().cache_editing_header_latched
1020}
1021
1022pub fn set_cache_editing_header_latched(v: bool) {
1023    STATE.lock().unwrap().cache_editing_header_latched = Some(v);
1024}
1025
1026pub fn get_thinking_clear_latched() -> Option<bool> {
1027    STATE.lock().unwrap().thinking_clear_latched
1028}
1029
1030pub fn set_thinking_clear_latched(v: bool) {
1031    STATE.lock().unwrap().thinking_clear_latched = Some(v);
1032}
1033
1034pub fn clear_beta_header_latches() {
1035    let mut state = STATE.lock().unwrap();
1036    state.afk_mode_header_latched = None;
1037    state.fast_mode_header_latched = None;
1038    state.cache_editing_header_latched = None;
1039    state.thinking_clear_latched = None;
1040}
1041
1042pub fn get_prompt_id() -> Option<String> {
1043    STATE.lock().unwrap().prompt_id.clone()
1044}
1045
1046pub fn set_prompt_id(id: Option<String>) {
1047    STATE.lock().unwrap().prompt_id = id;
1048}
1049
1050struct State {
1051    pub original_cwd: String,
1052    pub project_root: String,
1053    pub total_cost_usd: f64,
1054    pub total_api_duration: f64,
1055    pub total_api_duration_without_retries: f64,
1056    pub total_tool_duration: f64,
1057    pub turn_hook_duration_ms: f64,
1058    pub turn_tool_duration_ms: f64,
1059    pub turn_classifier_duration_ms: f64,
1060    pub turn_tool_count: u64,
1061    pub turn_hook_count: u64,
1062    pub turn_classifier_count: u64,
1063    pub start_time: u64,
1064    pub last_interaction_time: u64,
1065    pub total_lines_added: u64,
1066    pub total_lines_removed: u64,
1067    pub has_unknown_model_cost: bool,
1068    pub cwd: String,
1069    pub model_usage: HashMap<String, ModelUsage>,
1070    pub main_loop_model_override: Option<ModelSetting>,
1071    pub initial_main_loop_model: Option<ModelSetting>,
1072    pub model_strings: Option<ModelStrings>,
1073    pub is_interactive: bool,
1074    pub kairos_active: bool,
1075    pub strict_tool_result_pairing: bool,
1076    pub sdk_agent_progress_summaries_enabled: bool,
1077    pub user_msg_opt_in: bool,
1078    pub client_type: String,
1079    pub session_source: Option<String>,
1080    pub question_preview_format: Option<String>,
1081    pub flag_settings_path: Option<String>,
1082    pub flag_settings_inline: Option<HashMap<String, serde_json::Value>>,
1083    pub allowed_setting_sources: Vec<String>,
1084    pub session_ingress_token: Option<String>,
1085    pub oauth_token_from_fd: Option<String>,
1086    pub api_key_from_fd: Option<String>,
1087    pub stats_store: Option<()>,
1088    pub session_id: String,
1089    pub parent_session_id: Option<String>,
1090    pub agent_color_map: HashMap<String, String>,
1091    pub agent_color_index: usize,
1092    pub last_api_request: Option<serde_json::Value>,
1093    pub last_api_request_messages: Option<serde_json::Value>,
1094    pub last_classifier_requests: Option<Vec<serde_json::Value>>,
1095    pub cached_claude_md_content: Option<String>,
1096    pub in_memory_error_log: Vec<ErrorLogEntry>,
1097    pub inline_plugins: Vec<String>,
1098    pub chrome_flag_override: Option<bool>,
1099    pub use_cowork_plugins: bool,
1100    pub session_bypass_permissions_mode: bool,
1101    pub scheduled_tasks_enabled: bool,
1102    pub session_cron_tasks: Vec<SessionCronTask>,
1103    pub session_created_teams: HashSet<String>,
1104    pub session_trust_accepted: bool,
1105    pub session_persistence_disabled: bool,
1106    pub has_exited_plan_mode: bool,
1107    pub needs_plan_mode_exit_attachment: bool,
1108    pub needs_auto_mode_exit_attachment: bool,
1109    pub lsp_recommendation_shown_this_session: bool,
1110    pub init_json_schema: Option<HashMap<String, serde_json::Value>>,
1111    pub registered_hooks: Option<HashMap<String, Vec<HookMatcher>>>,
1112    pub plan_slug_cache: HashMap<String, String>,
1113    pub teleported_session_info: Option<TeleportedSessionInfo>,
1114    pub invoked_skills: HashMap<String, InvokedSkillInfo>,
1115    pub slow_operations: Vec<SlowOperation>,
1116    pub sdk_betas: Option<Vec<String>>,
1117    pub main_thread_agent_type: Option<String>,
1118    pub is_remote_mode: bool,
1119    pub direct_connect_server_url: Option<String>,
1120    pub system_prompt_section_cache: HashMap<String, Option<String>>,
1121    pub last_emitted_date: Option<String>,
1122    pub additional_directories_for_claude_md: Vec<String>,
1123    pub allowed_channels: Vec<ChannelEntry>,
1124    pub has_dev_channels: bool,
1125    pub session_project_dir: Option<String>,
1126    pub prompt_cache_1h_allowlist: Option<Vec<String>>,
1127    pub prompt_cache_1h_eligible: Option<bool>,
1128    pub afk_mode_header_latched: Option<bool>,
1129    pub fast_mode_header_latched: Option<bool>,
1130    pub cache_editing_header_latched: Option<bool>,
1131    pub thinking_clear_latched: Option<bool>,
1132    pub prompt_id: Option<String>,
1133    pub last_main_request_id: Option<String>,
1134    pub last_api_completion_timestamp: Option<u64>,
1135    pub pending_post_compaction: bool,
1136    /// Per-turn token snapshotting state
1137    pub output_tokens_at_turn_start: u64,
1138    pub current_turn_token_budget: Option<f64>,
1139    pub budget_continuation_count: u64,
1140}
1141
1142fn get_initial_state() -> State {
1143    let resolved_cwd = std::env::current_dir()
1144        .map(|p| p.to_string_lossy().to_string())
1145        .unwrap_or_else(|_| ".".to_string());
1146
1147    let session_id = Uuid::new_v4().to_string();
1148
1149    State {
1150        original_cwd: resolved_cwd.clone(),
1151        project_root: resolved_cwd.clone(),
1152        total_cost_usd: 0.0,
1153        total_api_duration: 0.0,
1154        total_api_duration_without_retries: 0.0,
1155        total_tool_duration: 0.0,
1156        turn_hook_duration_ms: 0.0,
1157        turn_tool_duration_ms: 0.0,
1158        turn_classifier_duration_ms: 0.0,
1159        turn_tool_count: 0,
1160        turn_hook_count: 0,
1161        turn_classifier_count: 0,
1162        start_time: current_timestamp(),
1163        last_interaction_time: current_timestamp(),
1164        total_lines_added: 0,
1165        total_lines_removed: 0,
1166        has_unknown_model_cost: false,
1167        cwd: resolved_cwd,
1168        model_usage: HashMap::new(),
1169        main_loop_model_override: None,
1170        initial_main_loop_model: None,
1171        model_strings: None,
1172        is_interactive: false,
1173        kairos_active: false,
1174        strict_tool_result_pairing: false,
1175        sdk_agent_progress_summaries_enabled: false,
1176        user_msg_opt_in: false,
1177        client_type: "cli".to_string(),
1178        session_source: None,
1179        question_preview_format: None,
1180        session_ingress_token: None,
1181        oauth_token_from_fd: None,
1182        api_key_from_fd: None,
1183        flag_settings_path: None,
1184        flag_settings_inline: None,
1185        allowed_setting_sources: vec![
1186            "userSettings".to_string(),
1187            "projectSettings".to_string(),
1188            "localSettings".to_string(),
1189            "flagSettings".to_string(),
1190            "policySettings".to_string(),
1191        ],
1192        stats_store: None,
1193        session_id,
1194        parent_session_id: None,
1195        agent_color_map: HashMap::new(),
1196        agent_color_index: 0,
1197        last_api_request: None,
1198        last_api_request_messages: None,
1199        last_classifier_requests: None,
1200        cached_claude_md_content: None,
1201        in_memory_error_log: Vec::new(),
1202        inline_plugins: Vec::new(),
1203        chrome_flag_override: None,
1204        use_cowork_plugins: false,
1205        session_bypass_permissions_mode: false,
1206        scheduled_tasks_enabled: false,
1207        session_cron_tasks: Vec::new(),
1208        session_created_teams: HashSet::new(),
1209        session_trust_accepted: false,
1210        session_persistence_disabled: false,
1211        has_exited_plan_mode: false,
1212        needs_plan_mode_exit_attachment: false,
1213        needs_auto_mode_exit_attachment: false,
1214        lsp_recommendation_shown_this_session: false,
1215        init_json_schema: None,
1216        registered_hooks: None,
1217        plan_slug_cache: HashMap::new(),
1218        teleported_session_info: None,
1219        invoked_skills: HashMap::new(),
1220        slow_operations: Vec::new(),
1221        sdk_betas: None,
1222        main_thread_agent_type: None,
1223        is_remote_mode: false,
1224        direct_connect_server_url: None,
1225        system_prompt_section_cache: HashMap::new(),
1226        last_emitted_date: None,
1227        additional_directories_for_claude_md: Vec::new(),
1228        allowed_channels: Vec::new(),
1229        has_dev_channels: false,
1230        session_project_dir: None,
1231        prompt_cache_1h_allowlist: None,
1232        prompt_cache_1h_eligible: None,
1233        afk_mode_header_latched: None,
1234        fast_mode_header_latched: None,
1235        cache_editing_header_latched: None,
1236        thinking_clear_latched: None,
1237        prompt_id: None,
1238        last_main_request_id: None,
1239        last_api_completion_timestamp: None,
1240        pending_post_compaction: false,
1241        output_tokens_at_turn_start: 0,
1242        current_turn_token_budget: None,
1243        budget_continuation_count: 0,
1244    }
1245}
1246
1247static STATE: Lazy<Mutex<State>> = Lazy::new(|| Mutex::new(get_initial_state()));
ai_agent/bootstrap/state.rs

ai_agent/bootstrap/
state.rs
