Skip to content

Commit d281af7

Browse files
feat: Adds experiment manager (#3054)
* feat: Adds experiment manager * chore: Includes experimental commands in prompts.rs --------- Co-authored-by: Kenneth S. <[email protected]>
1 parent 7e28138 commit d281af7

File tree

14 files changed

+468
-221
lines changed

14 files changed

+468
-221
lines changed

crates/chat-cli/src/cli/chat/cli/checkpoint.rs

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,10 @@ use crate::cli::chat::{
2323
ChatSession,
2424
ChatState,
2525
};
26-
use crate::database::settings::Setting;
26+
use crate::cli::experiment::experiment_manager::{
27+
ExperimentManager,
28+
ExperimentName,
29+
};
2730
use crate::os::Os;
2831
use crate::util::directories::get_shadow_repo_dir;
2932

@@ -84,12 +87,7 @@ With --hard:
8487
impl CheckpointSubcommand {
8588
pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
8689
// Check if checkpoint is enabled
87-
if !os
88-
.database
89-
.settings
90-
.get_bool(Setting::EnabledCheckpoint)
91-
.unwrap_or(false)
92-
{
90+
if !ExperimentManager::is_enabled(os, ExperimentName::Checkpoint) {
9391
execute!(
9492
session.stderr,
9593
style::SetForegroundColor(Color::Red),

crates/chat-cli/src/cli/chat/cli/experiment.rs

Lines changed: 14 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -17,50 +17,9 @@ use crate::cli::chat::{
1717
ChatSession,
1818
ChatState,
1919
};
20-
use crate::database::settings::Setting;
20+
use crate::cli::experiment::experiment_manager::ExperimentManager;
2121
use crate::os::Os;
2222

23-
/// Represents an experimental feature that can be toggled
24-
#[derive(Debug, Clone)]
25-
struct Experiment {
26-
name: &'static str,
27-
description: &'static str,
28-
setting_key: Setting,
29-
}
30-
31-
static AVAILABLE_EXPERIMENTS: &[Experiment] = &[
32-
Experiment {
33-
name: "Knowledge",
34-
description: "Enables persistent context storage and retrieval across chat sessions (/knowledge)",
35-
setting_key: Setting::EnabledKnowledge,
36-
},
37-
Experiment {
38-
name: "Thinking",
39-
description: "Enables complex reasoning with step-by-step thought processes",
40-
setting_key: Setting::EnabledThinking,
41-
},
42-
Experiment {
43-
name: "Tangent Mode",
44-
description: "Enables entering into a temporary mode for sending isolated conversations (/tangent)",
45-
setting_key: Setting::EnabledTangentMode,
46-
},
47-
Experiment {
48-
name: "Todo Lists",
49-
description: "Enables Q to create todo lists that can be viewed and managed using /todos",
50-
setting_key: Setting::EnabledTodoList,
51-
},
52-
Experiment {
53-
name: "Checkpoint",
54-
description: "Enables workspace checkpoints to snapshot, list, expand, diff, and restore files (/checkpoint)\nNote: Cannot be used in tangent mode (to avoid mixing up conversation history)",
55-
setting_key: Setting::EnabledCheckpoint,
56-
},
57-
Experiment {
58-
name: "Context Usage Indicator",
59-
description: "Shows context usage percentage in the prompt (e.g., [rust-agent] 6% >)",
60-
setting_key: Setting::EnabledContextUsageIndicator,
61-
},
62-
];
63-
6423
#[derive(Debug, PartialEq, Args)]
6524
pub struct ExperimentArgs;
6625
impl ExperimentArgs {
@@ -75,9 +34,10 @@ async fn select_experiment(os: &mut Os, session: &mut ChatSession) -> Result<Opt
7534
// Get current experiment status
7635
let mut experiment_labels = Vec::new();
7736
let mut current_states = Vec::new();
37+
let experiments = ExperimentManager::get_experiments();
7838

79-
for experiment in AVAILABLE_EXPERIMENTS {
80-
let is_enabled = os.database.settings.get_bool(experiment.setting_key).unwrap_or(false);
39+
for experiment in &experiments {
40+
let is_enabled = ExperimentManager::is_enabled(os, experiment.experiment_name);
8141

8242
current_states.push(is_enabled);
8343

@@ -92,7 +52,7 @@ async fn select_experiment(os: &mut Os, session: &mut ChatSession) -> Result<Opt
9252

9353
let label = format!(
9454
"{:<25} {:<6} {}",
95-
experiment.name,
55+
experiment.experiment_name.as_str(),
9656
status_indicator,
9757
style::Stylize::dark_grey(description)
9858
);
@@ -145,37 +105,30 @@ async fn select_experiment(os: &mut Os, session: &mut ChatSession) -> Result<Opt
145105
)?;
146106

147107
// Skip if user selected disclaimer or empty line (last 2 items)
148-
if index >= AVAILABLE_EXPERIMENTS.len() {
108+
if index >= experiments.len() {
149109
return Ok(Some(ChatState::PromptUser {
150110
skip_printing_tools: false,
151111
}));
152112
}
153113

154-
let experiment = &AVAILABLE_EXPERIMENTS[index];
114+
let experiment = &experiments[index];
155115
let current_state = current_states[index];
156116
let new_state = !current_state;
157117

158-
// Update the setting
159-
os.database
160-
.settings
161-
.set(experiment.setting_key, new_state)
162-
.await
163-
.map_err(|e| ChatError::Custom(format!("Failed to update experiment setting: {e}").into()))?;
164-
165-
// Reload built-in tools to reflect the experiment change while preserving MCP tools
166-
session
167-
.conversation
168-
.reload_builtin_tools(os, &mut session.stderr)
169-
.await
170-
.map_err(|e| ChatError::Custom(format!("Failed to update tool spec: {e}").into()))?;
118+
// Update the setting using ExperimentManager
119+
ExperimentManager::set_enabled(os, experiment.experiment_name, new_state, session).await?;
171120

172121
let status_text = if new_state { "enabled" } else { "disabled" };
173122

174123
queue!(
175124
session.stderr,
176125
style::Print("\n"),
177126
style::SetForegroundColor(Color::Green),
178-
style::Print(format!(" {} experiment {}\n\n", experiment.name, status_text)),
127+
style::Print(format!(
128+
" {} experiment {}\n\n",
129+
experiment.experiment_name.as_str(),
130+
status_text
131+
)),
179132
style::ResetColor,
180133
style::SetForegroundColor(Color::Reset),
181134
style::SetBackgroundColor(Color::Reset),

crates/chat-cli/src/cli/chat/cli/knowledge.rs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,10 @@ use crate::cli::chat::{
1818
ChatSession,
1919
ChatState,
2020
};
21-
use crate::database::settings::Setting;
21+
use crate::cli::experiment::experiment_manager::{
22+
ExperimentManager,
23+
ExperimentName,
24+
};
2225
use crate::os::Os;
2326
use crate::util::knowledge_store::KnowledgeStore;
2427

@@ -79,10 +82,7 @@ impl KnowledgeSubcommand {
7982
}
8083

8184
fn is_feature_enabled(os: &Os) -> bool {
82-
os.database
83-
.settings
84-
.get_bool(Setting::EnabledKnowledge)
85-
.unwrap_or(false)
85+
ExperimentManager::is_enabled(os, ExperimentName::Knowledge)
8686
}
8787

8888
fn write_feature_disabled_message(session: &mut ChatSession) -> Result<(), std::io::Error> {

crates/chat-cli/src/cli/chat/cli/tangent.rs

Lines changed: 7 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,10 @@ use crate::cli::chat::{
1313
ChatSession,
1414
ChatState,
1515
};
16-
use crate::database::settings::Setting;
16+
use crate::cli::experiment::experiment_manager::{
17+
ExperimentManager,
18+
ExperimentName,
19+
};
1720
use crate::os::Os;
1821

1922
#[derive(Debug, PartialEq, Args)]
@@ -46,12 +49,7 @@ impl TangentArgs {
4649

4750
pub async fn execute(self, os: &Os, session: &mut ChatSession) -> Result<ChatState, ChatError> {
4851
// Check if tangent mode is enabled
49-
if !os
50-
.database
51-
.settings
52-
.get_bool(Setting::EnabledTangentMode)
53-
.unwrap_or(false)
54-
{
52+
if !ExperimentManager::is_enabled(os, ExperimentName::TangentMode) {
5553
execute!(
5654
session.stderr,
5755
style::SetForegroundColor(Color::Red),
@@ -66,12 +64,7 @@ impl TangentArgs {
6664
match self.subcommand {
6765
Some(TangentSubcommand::Tail) => {
6866
// Check if checkpoint is enabled
69-
if os
70-
.database
71-
.settings
72-
.get_bool(Setting::EnabledCheckpoint)
73-
.unwrap_or(false)
74-
{
67+
if ExperimentManager::is_enabled(os, ExperimentName::Checkpoint) {
7568
execute!(
7669
session.stderr,
7770
style::SetForegroundColor(Color::Yellow),
@@ -123,12 +116,7 @@ impl TangentArgs {
123116
)?;
124117
} else {
125118
// Check if checkpoint is enabled
126-
if os
127-
.database
128-
.settings
129-
.get_bool(Setting::EnabledCheckpoint)
130-
.unwrap_or(false)
131-
{
119+
if ExperimentManager::is_enabled(os, ExperimentName::Checkpoint) {
132120
execute!(
133121
session.stderr,
134122
style::SetForegroundColor(Color::Yellow),

crates/chat-cli/src/cli/chat/mod.rs

Lines changed: 21 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,10 @@ use crate::cli::chat::cli::prompts::{
156156
};
157157
use crate::cli::chat::message::UserMessage;
158158
use crate::cli::chat::util::sanitize_unicode_tags;
159+
use crate::cli::experiment::experiment_manager::{
160+
ExperimentManager,
161+
ExperimentName,
162+
};
159163
use crate::database::settings::Setting;
160164
use crate::os::Os;
161165
use crate::telemetry::core::{
@@ -1151,6 +1155,14 @@ impl ChatSession {
11511155

11521156
Ok(())
11531157
}
1158+
1159+
/// Reload built-in tools to reflect experiment changes while preserving MCP tools
1160+
pub async fn reload_builtin_tools(&mut self, os: &mut Os) -> Result<(), ChatError> {
1161+
self.conversation
1162+
.reload_builtin_tools(os, &mut self.stderr)
1163+
.await
1164+
.map_err(|e| ChatError::Custom(format!("Failed to update tool spec: {e}").into()))
1165+
}
11541166
}
11551167

11561168
impl Drop for ChatSession {
@@ -1331,12 +1343,7 @@ impl ChatSession {
13311343
}
13321344

13331345
// Initialize capturing if possible
1334-
if os
1335-
.database
1336-
.settings
1337-
.get_bool(Setting::EnabledCheckpoint)
1338-
.unwrap_or(false)
1339-
{
1346+
if ExperimentManager::is_enabled(os, ExperimentName::Checkpoint) {
13401347
let path = get_shadow_repo_dir(os, self.conversation.conversation_id().to_string())?;
13411348
let start = std::time::Instant::now();
13421349
let checkpoint_manager = match CheckpointManager::auto_init(os, &path, self.conversation.history()).await {
@@ -2125,12 +2132,7 @@ impl ChatSession {
21252132
} else {
21262133
// Track the message for checkpoint descriptions, but only if not already set
21272134
// This prevents tool approval responses (y/n/t) from overwriting the original message
2128-
if os
2129-
.database
2130-
.settings
2131-
.get_bool(Setting::EnabledCheckpoint)
2132-
.unwrap_or(false)
2133-
&& !self.conversation.is_in_tangent_mode()
2135+
if ExperimentManager::is_enabled(os, ExperimentName::Checkpoint) && !self.conversation.is_in_tangent_mode()
21342136
{
21352137
if let Some(manager) = self.conversation.checkpoint_manager.as_mut() {
21362138
if !manager.message_locked && self.pending_tool_index.is_none() {
@@ -2220,11 +2222,7 @@ impl ChatSession {
22202222

22212223
async fn tool_use_execute(&mut self, os: &mut Os) -> Result<ChatState, ChatError> {
22222224
// Check if we should auto-enter tangent mode for introspect tool
2223-
if os
2224-
.database
2225-
.settings
2226-
.get_bool(Setting::EnabledTangentMode)
2227-
.unwrap_or(false)
2225+
if ExperimentManager::is_enabled(os, ExperimentName::TangentMode)
22282226
&& os
22292227
.database
22302228
.settings
@@ -2365,13 +2363,10 @@ impl ChatSession {
23652363

23662364
// Handle checkpoint after tool execution - store tag for later display
23672365
let checkpoint_tag: Option<String> = {
2368-
let enabled = os
2369-
.database
2370-
.settings
2371-
.get_bool(Setting::EnabledCheckpoint)
2372-
.unwrap_or(false)
2373-
&& !self.conversation.is_in_tangent_mode();
2374-
if invoke_result.is_err() || !enabled {
2366+
if invoke_result.is_err()
2367+
|| !ExperimentManager::is_enabled(os, ExperimentName::Checkpoint)
2368+
|| self.conversation.is_in_tangent_mode()
2369+
{
23752370
None
23762371
}
23772372
// Take manager out temporarily to avoid borrow conflicts
@@ -2984,12 +2979,7 @@ impl ChatSession {
29842979
self.tool_turn_start_time = None;
29852980

29862981
// Create turn checkpoint if tools were used
2987-
if os
2988-
.database
2989-
.settings
2990-
.get_bool(Setting::EnabledCheckpoint)
2991-
.unwrap_or(false)
2992-
&& !self.conversation.is_in_tangent_mode()
2982+
if ExperimentManager::is_enabled(os, ExperimentName::Checkpoint) && !self.conversation.is_in_tangent_mode()
29932983
{
29942984
if let Some(mut manager) = self.conversation.checkpoint_manager.take() {
29952985
if manager.tools_in_turn > 0 {
@@ -3364,12 +3354,7 @@ impl ChatSession {
33643354
let tangent_mode = self.conversation.is_in_tangent_mode();
33653355

33663356
// Check if context usage indicator is enabled
3367-
let usage_percentage = if os
3368-
.database
3369-
.settings
3370-
.get_bool(crate::database::settings::Setting::EnabledContextUsageIndicator)
3371-
.unwrap_or(false)
3372-
{
3357+
let usage_percentage = if ExperimentManager::is_enabled(os, ExperimentName::ContextUsageIndicator) {
33733358
use crate::cli::chat::cli::usage::get_total_usage_percentage;
33743359
get_total_usage_percentage(self, os).await.ok()
33753360
} else {

0 commit comments

Comments
 (0)