(draft)feat: add cost estimation per provider for Goose CLI

.devcontainer/Dockerfile

@@ -7,4 +7,5 @@ RUN apt-get update && apt-get install -y \
     libdbus-1-dev \
     gnome-keyring \
     libxcb1-dev \
+    protobuf-compiler \
     && apt-get clean

Cargo.toml

@@ -4,7 +4,7 @@ resolver = "2"
 
 [workspace.package]
 edition = "2021"
-version = "1.0.32"
+version = "1.0.35"
 authors = ["Block <[email protected]>"]
 license = "Apache-2.0"
 repository = "https://github.com/block/goose"

crates/goose-cli/src/commands/configure.rs

@@ -498,8 +498,13 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
         )
         .item(
             "sse",
-            "Remote Extension",
-            "Connect to a remote extension via SSE",
+            "Remote Extension (SSE)",
+            "Connect to a remote extension via Server-Sent Events",
+        )
+        .item(
+            "streamable_http",
+            "Remote Extension (Streaming HTTP)",
+            "Connect to a remote extension via MCP Streaming HTTP",
         )
         .interact()?;
 
@@ -767,6 +772,133 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
 
             cliclack::outro(format!("Added {} extension", style(name).green()))?;
         }
+        "streamable_http" => {
+            let extensions = ExtensionConfigManager::get_all_names()?;
+            let name: String = cliclack::input("What would you like to call this extension?")
+                .placeholder("my-remote-extension")
+                .validate(move |input: &String| {
+                    if input.is_empty() {
+                        Err("Please enter a name")
+                    } else if extensions.contains(input) {
+                        Err("An extension with this name already exists")
+                    } else {
+                        Ok(())
+                    }
+                })
+                .interact()?;
+
+            let uri: String = cliclack::input("What is the Streaming HTTP endpoint URI?")
+                .placeholder("http://localhost:8000/messages")
+                .validate(|input: &String| {
+                    if input.is_empty() {
+                        Err("Please enter a URI")
+                    } else if !(input.starts_with("http://") || input.starts_with("https://")) {
+                        Err("URI should start with http:// or https://")
+                    } else {
+                        Ok(())
+                    }
+                })
+                .interact()?;
+
+            let timeout: u64 = cliclack::input("Please set the timeout for this tool (in secs):")
+                .placeholder(&goose::config::DEFAULT_EXTENSION_TIMEOUT.to_string())
+                .validate(|input: &String| match input.parse::<u64>() {
+                    Ok(_) => Ok(()),
+                    Err(_) => Err("Please enter a valid timeout"),
+                })
+                .interact()?;
+
+            let add_desc = cliclack::confirm("Would you like to add a description?").interact()?;
+
+            let description = if add_desc {
+                let desc = cliclack::input("Enter a description for this extension:")
+                    .placeholder("Description")
+                    .validate(|input: &String| {
+                        if input.trim().is_empty() {
+                            Err("Please enter a valid description")
+                        } else {
+                            Ok(())
+                        }
+                    })
+                    .interact()?;
+                Some(desc)
+            } else {
+                None
+            };
+
+            let add_headers =
+                cliclack::confirm("Would you like to add custom headers?").interact()?;
+
+            let mut headers = HashMap::new();
+            if add_headers {
+                loop {
+                    let key: String = cliclack::input("Header name:")
+                        .placeholder("Authorization")
+                        .interact()?;
+
+                    let value: String = cliclack::input("Header value:")
+                        .placeholder("Bearer token123")
+                        .interact()?;
+
+                    headers.insert(key, value);
+
+                    if !cliclack::confirm("Add another header?").interact()? {
+                        break;
+                    }
+                }
+            }
+
+            let add_env = false; // No env prompt for Streaming HTTP
+
+            let mut envs = HashMap::new();
+            let mut env_keys = Vec::new();
+            let config = Config::global();
+
+            if add_env {
+                loop {
+                    let key: String = cliclack::input("Environment variable name:")
+                        .placeholder("API_KEY")
+                        .interact()?;
+
+                    let value: String = cliclack::password("Environment variable value:")
+                        .mask('▪')
+                        .interact()?;
+
+                    // Try to store in keychain
+                    let keychain_key = key.to_string();
+                    match config.set_secret(&keychain_key, Value::String(value.clone())) {
+                        Ok(_) => {
+                            // Successfully stored in keychain, add to env_keys
+                            env_keys.push(keychain_key);
+                        }
+                        Err(_) => {
+                            // Failed to store in keychain, store directly in envs
+                            envs.insert(key, value);
+                        }
+                    }
+
+                    if !cliclack::confirm("Add another environment variable?").interact()? {
+                        break;
+                    }
+                }
+            }
+
+            ExtensionConfigManager::set(ExtensionEntry {
+                enabled: true,
+                config: ExtensionConfig::StreamableHttp {
+                    name: name.clone(),
+                    uri,
+                    envs: Envs::new(envs),
+                    env_keys,
+                    headers,
+                    description,
+                    timeout: Some(timeout),
+                    bundled: None,
+                },
+            })?;
+
+            cliclack::outro(format!("Added {} extension", style(name).green()))?;
+        }
         _ => unreachable!(),
     };
 

crates/goose-cli/src/logging.rs

@@ -224,23 +224,26 @@ mod tests {
     }
 
     async fn do_test_log_file_name(session_name: Option<&str>, _with_error_capture: bool) {
-        // Create a unique test directory for each test
-        let test_name = session_name.unwrap_or("no_session");
-        let random_suffix = rand::random::<u32>() % 100000000;
-        let test_dir = PathBuf::from(format!(
-            "/tmp/goose_test_home_{}_{}",
-            test_name, random_suffix
-        ));
-        if test_dir.exists() {
-            fs::remove_dir_all(&test_dir).unwrap();
-        }
-        fs::create_dir_all(&test_dir).unwrap();
+        use tempfile::TempDir;
+
+        // Create a unique prefix to avoid test interference
+        let test_id = format!(
+            "{}_{}",
+            session_name.unwrap_or("no_session"),
+            rand::random::<u32>()
+        );
+
+        // Create a proper temporary directory that will be automatically cleaned up
+        let _temp_dir = TempDir::with_prefix(&format!("goose_test_{}_", test_id)).unwrap();
+        let test_dir = _temp_dir.path();
 
         // Set up environment
         if cfg!(windows) {
-            env::set_var("USERPROFILE", &test_dir);
+            env::set_var("USERPROFILE", test_dir);
         } else {
-            env::set_var("HOME", &test_dir);
+            env::set_var("HOME", test_dir);
+            // Also set TMPDIR to prevent temp directory sharing between tests
+            env::set_var("TMPDIR", test_dir);
         }
 
         // Create error capture if needed - but don't use it in tests to avoid tokio runtime issues
@@ -251,8 +254,10 @@ mod tests {
         println!("Before timestamp: {}", before_timestamp);
 
         // Get the log directory and clean any existing log files
+        let random_suffix = rand::random::<u32>() % 100000000;
         let log_dir = get_log_directory_with_date(Some(format!("test-{}", random_suffix))).unwrap();
         println!("Log directory: {}", log_dir.display());
+        println!("Test directory: {}", test_dir.display());
         if log_dir.exists() {
             for entry in fs::read_dir(&log_dir).unwrap() {
                 let entry = entry.unwrap();
@@ -429,10 +434,8 @@ mod tests {
         // Wait a moment to ensure all files are written
         std::thread::sleep(std::time::Duration::from_millis(100));
 
-        // Clean up test directory
-        fs::remove_dir_all(&test_dir).unwrap_or_else(|e| {
-            println!("Warning: Failed to clean up test directory: {}", e);
-        });
+        // Keep _temp_dir alive until the end so it doesn't get cleaned up prematurely
+        drop(_temp_dir);
     }
 
     #[tokio::test]

crates/goose-cli/src/session/mod.rs

@@ -23,6 +23,7 @@ use goose::agents::extension::{Envs, ExtensionConfig};
 use goose::agents::{Agent, SessionConfig};
 use goose::config::Config;
 use goose::message::{Message, MessageContent};
+use goose::providers::pricing::initialize_pricing_cache;
 use goose::session;
 use input::InputResult;
 use mcp_core::handler::ToolError;
@@ -1303,13 +1304,42 @@ impl Session {
     pub async fn display_context_usage(&self) -> Result<()> {
         let provider = self.agent.provider().await?;
         let model_config = provider.get_model_config();
-        let context_limit = model_config.context_limit.unwrap_or(32000);
+        let context_limit = model_config.context_limit();
+
+        let config = Config::global();
+        let show_cost = config
+            .get_param::<bool>("GOOSE_CLI_SHOW_COST")
+            .unwrap_or(false);
+
+        let provider_name = config
+            .get_param::<String>("GOOSE_PROVIDER")
+            .unwrap_or_else(|_| "unknown".to_string());
+
+        // Initialize pricing cache on startup
+        tracing::info!("Initializing pricing cache...");
+        if let Err(e) = initialize_pricing_cache().await {
+            tracing::warn!(
+                "Failed to initialize pricing cache: {e}. Pricing data may not be available."
+            );
+        }
 
         match self.get_metadata() {
             Ok(metadata) => {
                 let total_tokens = metadata.total_tokens.unwrap_or(0) as usize;
 
                 output::display_context_usage(total_tokens, context_limit);
+
+                if show_cost {
+                    let input_tokens = metadata.input_tokens.unwrap_or(0) as usize;
+                    let output_tokens = metadata.output_tokens.unwrap_or(0) as usize;
+                    output::display_cost_usage(
+                        &provider_name,
+                        &model_config.model_name,
+                        input_tokens,
+                        output_tokens,
+                    )
+                    .await;
+                }
             }
             Err(_) => {
                 output::display_context_usage(0, context_limit);
@@ -1450,7 +1480,8 @@ fn get_reasoner() -> Result<Arc<dyn Provider>, anyhow::Error> {
             .expect("No model configured. Run 'goose configure' first")
     };
 
-    let model_config = ModelConfig::new(model);
+    let model_config =
+        ModelConfig::new_with_context_env(model, Some("GOOSE_PLANNER_CONTEXT_LIMIT"));
     let reasoner = create(&provider, model_config)?;
 
     Ok(reasoner)