From 25820bc7454dca11e55484c98d5799b8927676d4 Mon Sep 17 00:00:00 2001 From: Oded Falik Date: Tue, 13 Jan 2026 14:19:23 -0800 Subject: [PATCH] feat: lazy API key loading for MCP server Add support for deferring API key command execution until first use: - MCP server now defaults to lazy loading (defers LAMBDA_API_KEY_COMMAND execution until the first API request) - Add --eager flag to MCP server for startup-time API key loading - API key is cached in memory after first retrieval - CLI continues to load API key at startup (unchanged behavior) This improves Claude Code startup time when using slow API key commands (e.g., 1Password) since the command only runs when Lambda tools are actually used. Co-Authored-By: Claude Opus 4.5 --- README.md | 19 ++++++ src/api.rs | 134 +++++++++++++++++++++++++++++++----------- src/bin/lambda-mcp.rs | 14 +++-- 3 files changed, 129 insertions(+), 38 deletions(-) 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);