diff --git a/README.md b/README.md index 896cc45..6887914 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,7 @@ [![npm](https://img.shields.io/npm/v/@strand-ai/lambda-mcp)](https://www.npmjs.com/package/@strand-ai/lambda-mcp) [![MCP](https://img.shields.io/badge/MCP-compatible-8A2BE2)](https://modelcontextprotocol.io) [![Install in VS Code](https://img.shields.io/badge/VS_Code-Install_Server-0098FF?logo=visualstudiocode&logoColor=white)](https://vscode.dev/redirect/mcp/install?name=lambda-labs&config=%7B%22command%22%3A%22npx%22%2C%22args%22%3A%5B%22-y%22%2C%22%40strand-ai%2Flambda-mcp%22%5D%7D) +[![Install in Cursor](https://img.shields.io/badge/Cursor-Install_Server-000000?logo=cursor&logoColor=white)](cursor://anysphere.cursor-deeplink/mcp/install?name=lambda-labs&config=%7B%22command%22%3A%22npx%22%2C%22args%22%3A%5B%22-y%22%2C%22%40strand-ai%2Flambda-mcp%22%5D%7D) A fast CLI and MCP server for managing [Lambda Labs](https://lambdalabs.com/) cloud GPU instances. @@ -111,6 +112,24 @@ The easiest way to use `lambda-mcp` is via npx—no installation required: npx @strand-ai/lambda-mcp ``` +### Options + +| Flag | Description | +|------|-------------| +| `--eager` | Execute API key command at startup instead of on first use | + +#### API Key Loading + +When using `LAMBDA_API_KEY_COMMAND`, the MCP server defers command execution until the first API request by default. This avoids unnecessary delays when starting Claude Code if you don't use Lambda tools in every session. + +Use `--eager` to execute the command at startup instead: + +```bash +npx @strand-ai/lambda-mcp --eager +``` + +> **Note:** The CLI (`lambda`) always executes the API key command at startup since it's used for immediate operations. + ### Available Tools | Tool | Description | diff --git a/src/api.rs b/src/api.rs index 69a744e..d99d9e9 100644 --- a/src/api.rs +++ b/src/api.rs @@ -3,6 +3,7 @@ use reqwest::header::AUTHORIZATION; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +use std::sync::Mutex; use std::time::Duration; use thiserror::Error; @@ -106,10 +107,21 @@ pub struct Region { pub description: String, } +/// Source for the API key - either a direct value or a command to execute +#[derive(Debug, Clone)] +enum ApiKeySource { + /// Direct API key value (already resolved) + Direct(String), + /// Command to execute to get the API key (lazy evaluation) + Command(String), +} + /// Lambda Labs API client pub struct LambdaClient { client: Client, - api_key: String, + api_key_source: ApiKeySource, + /// Cached API key (used for lazy evaluation) + cached_api_key: Mutex>, } impl LambdaClient { @@ -120,7 +132,26 @@ impl LambdaClient { .build() .context("Failed to create HTTP client")?; - Ok(Self { client, api_key }) + Ok(Self { + client, + api_key_source: ApiKeySource::Direct(api_key), + cached_api_key: Mutex::new(None), + }) + } + + /// Create a client with a lazy API key source (command executed on first use) + fn new_lazy(api_key_source: ApiKeySource) -> Result { + let client = Client::builder() + .timeout(Duration::from_secs(DEFAULT_TIMEOUT_SECS)) + .connect_timeout(Duration::from_secs(10)) + .build() + .context("Failed to create HTTP client")?; + + Ok(Self { + client, + api_key_source, + cached_api_key: Mutex::new(None), + }) } /// Create a client using environment variables for the API key. @@ -128,18 +159,70 @@ impl LambdaClient { /// Checks in order: /// 1. `LAMBDA_API_KEY` - Direct API key /// 2. `LAMBDA_API_KEY_COMMAND` - Command to execute to get the API key (e.g., `op read op://vault/lambda/api-key`) + /// + /// By default, if `LAMBDA_API_KEY_COMMAND` is used, the command is executed immediately. pub fn from_env() -> Result { - let api_key = get_api_key_from_env()?; - Self::new(api_key) + Self::from_env_with_options(false) + } + + /// Create a client using environment variables for the API key with options. + /// + /// If `lazy` is true and `LAMBDA_API_KEY_COMMAND` is used, the command execution + /// is deferred until the first API request. + pub fn from_env_with_options(lazy: bool) -> Result { + // First, try direct API key (always immediate) + if let Ok(key) = std::env::var("LAMBDA_API_KEY") { + if !key.is_empty() { + return Self::new(key); + } + } + + // Then, try command-based retrieval + if let Ok(command) = std::env::var("LAMBDA_API_KEY_COMMAND") { + if !command.is_empty() { + if lazy { + // Defer command execution until first API request + return Self::new_lazy(ApiKeySource::Command(command)); + } else { + // Execute command immediately (default behavior) + let key = execute_api_key_command(&command)?; + return Self::new(key); + } + } + } + + Err(LambdaError::ApiKeyNotSet.into()) + } + + /// Get the API key, executing the command if necessary (lazy evaluation) + fn get_api_key(&self) -> Result { + match &self.api_key_source { + ApiKeySource::Direct(key) => Ok(key.clone()), + ApiKeySource::Command(cmd) => { + let mut cache = self + .cached_api_key + .lock() + .map_err(|e| anyhow!("Failed to acquire lock: {}", e))?; + + if let Some(key) = cache.as_ref() { + return Ok(key.clone()); + } + + let key = execute_api_key_command(cmd)?; + *cache = Some(key.clone()); + Ok(key) + } + } } /// Validate the API key by making a test request pub async fn validate_api_key(&self) -> Result<()> { + let api_key = self.get_api_key()?; let url = format!("{}/instances", API_BASE_URL); let response = self .client .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .header(AUTHORIZATION, format!("Bearer {}", api_key)) .send() .await .context("Failed to connect to Lambda Labs API")?; @@ -159,11 +242,12 @@ impl LambdaClient { /// List all available instance types pub async fn list_instance_types(&self) -> Result> { + let api_key = self.get_api_key()?; let url = format!("{}/instance-types", API_BASE_URL); let response = self .client .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .header(AUTHORIZATION, format!("Bearer {}", api_key)) .send() .await .context("Failed to fetch instance types")?; @@ -202,11 +286,12 @@ impl LambdaClient { /// Get instance type details (for checking availability) pub async fn get_instance_type(&self, gpu: &str) -> Result> { + let api_key = self.get_api_key()?; let url = format!("{}/instance-types", API_BASE_URL); let response = self .client .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .header(AUTHORIZATION, format!("Bearer {}", api_key)) .send() .await .context("Failed to fetch instance types")?; @@ -280,10 +365,11 @@ impl LambdaClient { payload["name"] = serde_json::Value::String(instance_name.to_string()); } + let api_key = self.get_api_key()?; let response = self .client .post(&url) - .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .header(AUTHORIZATION, format!("Bearer {}", api_key)) .json(&payload) .send() .await @@ -314,6 +400,7 @@ impl LambdaClient { /// Terminate an instance pub async fn terminate_instance(&self, instance_id: &str) -> Result<()> { + let api_key = self.get_api_key()?; let url = format!("{}/instance-operations/terminate", API_BASE_URL); let payload = serde_json::json!({ "instance_ids": [instance_id] @@ -322,7 +409,7 @@ impl LambdaClient { let response = self .client .post(&url) - .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .header(AUTHORIZATION, format!("Bearer {}", api_key)) .json(&payload) .send() .await @@ -338,11 +425,12 @@ impl LambdaClient { /// List all running instances pub async fn list_running_instances(&self) -> Result> { + let api_key = self.get_api_key()?; let url = format!("{}/instances", API_BASE_URL); let response = self .client .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .header(AUTHORIZATION, format!("Bearer {}", api_key)) .send() .await .context("Failed to fetch running instances")?; @@ -362,11 +450,12 @@ impl LambdaClient { /// Get details for a specific instance pub async fn get_instance(&self, instance_id: &str) -> Result { + let api_key = self.get_api_key()?; let url = format!("{}/instances/{}", API_BASE_URL, instance_id); let response = self .client .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .header(AUTHORIZATION, format!("Bearer {}", api_key)) .send() .await .context("Failed to fetch instance details")?; @@ -418,29 +507,6 @@ pub struct LaunchResult { pub region: String, } -/// Get API key from environment, supporting both direct key and command-based retrieval. -/// -/// Checks in order: -/// 1. `LAMBDA_API_KEY` - Direct API key value -/// 2. `LAMBDA_API_KEY_COMMAND` - Shell command to execute (e.g., `op read op://vault/lambda/api-key`) -fn get_api_key_from_env() -> Result { - // First, try direct API key - if let Ok(key) = std::env::var("LAMBDA_API_KEY") { - if !key.is_empty() { - return Ok(key); - } - } - - // Then, try command-based retrieval - if let Ok(command) = std::env::var("LAMBDA_API_KEY_COMMAND") { - if !command.is_empty() { - return execute_api_key_command(&command); - } - } - - Err(LambdaError::ApiKeyNotSet.into()) -} - /// Execute a shell command to retrieve the API key. fn execute_api_key_command(command: &str) -> Result { use std::process::Command; diff --git a/src/bin/lambda-mcp.rs b/src/bin/lambda-mcp.rs index 47a4135..965b74c 100644 --- a/src/bin/lambda-mcp.rs +++ b/src/bin/lambda-mcp.rs @@ -5,7 +5,7 @@ use rmcp::handler::server::wrapper::Parameters; use rmcp::model::{CallToolResult, Content, ServerCapabilities, ServerInfo}; use rmcp::schemars::JsonSchema; use rmcp::serde::Deserialize; -use rmcp::{tool, tool_router, ErrorData as McpError, ServerHandler, ServiceExt}; +use rmcp::{tool, tool_handler, tool_router, ErrorData as McpError, ServerHandler, ServiceExt}; use std::sync::Arc; /// Lambda Labs MCP Server @@ -17,9 +17,9 @@ struct LambdaService { } impl LambdaService { - fn new() -> Result { + fn new(lazy: bool) -> Result { dotenv::dotenv().ok(); - let client = LambdaClient::from_env()?; + let client = LambdaClient::from_env_with_options(lazy)?; Ok(Self { client: Arc::new(client), tool_router: Self::tool_router(), @@ -203,6 +203,7 @@ impl LambdaService { } } +#[tool_handler] impl ServerHandler for LambdaService { fn get_info(&self) -> ServerInfo { ServerInfo { @@ -220,8 +221,13 @@ impl ServerHandler for LambdaService { #[tokio::main] async fn main() -> Result<()> { + // Parse command line arguments + let args: Vec = std::env::args().collect(); + // Lazy loading is the default for MCP servers; use --eager to load API key at startup + let lazy = !args.iter().any(|arg| arg == "--eager"); + // Initialize the service - let service = match LambdaService::new() { + let service = match LambdaService::new(lazy) { Ok(s) => s, Err(e) => { eprintln!("Failed to initialize Lambda service: {}", e);