diff --git a/.github/workflows/release-please.yml b/.github/workflows/release-please.yml index f5c1bad..a541979 100644 --- a/.github/workflows/release-please.yml +++ b/.github/workflows/release-please.yml @@ -201,7 +201,7 @@ jobs: run: | cat > homebrew-tap/Formula/lambda-cli.rb << 'EOF' class LambdaCli < Formula - desc "CLI and MCP server for Lambda Labs cloud GPU instances" + desc "Unofficial CLI and MCP server for Lambda cloud GPU instances" homepage "https://github.com/Strand-AI/lambda-cli" url "https://github.com/Strand-AI/lambda-cli/archive/refs/tags/v${{ needs.release.outputs.version }}.tar.gz" sha256 "${{ steps.sha256.outputs.sha256 }}" diff --git a/Cargo.toml b/Cargo.toml index 766675e..78ee32e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ name = "lambda_cli" version = "0.4.0" edition = "2021" -description = "CLI and MCP server for Lambda Labs cloud GPU instances" +description = "Unofficial CLI and MCP server for Lambda cloud GPU instances" license = "MIT" repository = "https://github.com/Strand-AI/lambda-cli" diff --git a/README.md b/README.md index 8093242..14cb771 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,7 @@ # Lambda CLI +> **⚠️ UNOFFICIAL PROJECT** — This is a community-built tool, not affiliated with or endorsed by Lambda. + [![CI](https://github.com/Strand-AI/lambda-cli/actions/workflows/ci.yml/badge.svg)](https://github.com/Strand-AI/lambda-cli/actions/workflows/ci.yml) [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE) [![npm](https://img.shields.io/npm/v/@strand-ai/lambda-mcp)](https://www.npmjs.com/package/@strand-ai/lambda-mcp) @@ -7,7 +9,7 @@ [![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. +A fast CLI and MCP server for managing [Lambda](https://lambdalabs.com/) cloud GPU instances. **Two ways to use it:** - **CLI** (`lambda`) - Direct terminal commands for managing GPU instances @@ -30,7 +32,7 @@ Download from [GitHub Releases](https://github.com/Strand-AI/lambda-cli/releases ## Authentication -Get your API key from the [Lambda Labs dashboard](https://cloud.lambdalabs.com/api-keys). +Get your API key from the [Lambda dashboard](https://cloud.lambdalabs.com/api-keys). ### Option 1: Environment Variable ```bash @@ -44,6 +46,36 @@ export LAMBDA_API_KEY_COMMAND="op read op://Personal/Lambda/api-key" The command is executed at startup and its output is used as the API key. This works with any secret manager. +## Notifications (Optional) + +Get notified on Slack, Discord, or Telegram when your instance is ready and SSH-able. + +### Configuration + +Set one or more of these environment variables: + +```bash +# Slack (incoming webhook) +export LAMBDA_NOTIFY_SLACK_WEBHOOK="https://hooks.slack.com/services/T00/B00/XXX" + +# Discord (webhook URL) +export LAMBDA_NOTIFY_DISCORD_WEBHOOK="https://discord.com/api/webhooks/123/abc" + +# Telegram (bot token + chat ID) +export LAMBDA_NOTIFY_TELEGRAM_BOT_TOKEN="123456:ABC-DEF..." +export LAMBDA_NOTIFY_TELEGRAM_CHAT_ID="123456789" +``` + +### Setup Guides + +**Slack:** Create an [Incoming Webhook](https://api.slack.com/messaging/webhooks) in your workspace. + +**Discord:** In channel settings → Integrations → Webhooks → New Webhook → Copy Webhook URL. + +**Telegram:** +1. Message [@BotFather](https://t.me/botfather) → `/newbot` → copy the token +2. Message your bot, then visit `https://api.telegram.org/bot/getUpdates` to find your chat ID + --- ## CLI Usage @@ -89,6 +121,8 @@ lambda find --gpu gpu_8x_h100 --ssh my-key --interval 30 | `-s, --ssh` | SSH key name (required) | | `-n, --name` | Instance name | | `-r, --region` | Region (auto-selects if omitted) | +| `-f, --filesystem` | Filesystem to attach (must be in same region) | +| `--no-notify` | Disable notifications even if env vars are set | #### find | Flag | Description | @@ -97,12 +131,19 @@ lambda find --gpu gpu_8x_h100 --ssh my-key --interval 30 | `-s, --ssh` | SSH key name (required) | | `--interval` | Poll interval in seconds (default: 10) | | `-n, --name` | Instance name when launched | +| `-f, --filesystem` | Filesystem to attach when launched | +| `--no-notify` | Disable notifications even if env vars are set | + +Notifications are **automatic** when env vars are configured. Use `--no-notify` to disable: +```bash +lambda start --gpu gpu_1x_a10 --ssh my-key --no-notify +``` --- ## MCP Server -The `lambda-mcp` binary is an [MCP (Model Context Protocol)](https://modelcontextprotocol.io/) server that lets AI assistants manage your Lambda Labs infrastructure. +The `lambda-mcp` binary is an [MCP (Model Context Protocol)](https://modelcontextprotocol.io/) server that lets AI assistants manage your Lambda infrastructure. ### Quick Start with npx @@ -135,11 +176,15 @@ npx @strand-ai/lambda-mcp --eager | Tool | Description | |------|-------------| | `list_gpu_types` | List all available GPU instance types with pricing, specs, and current availability | -| `start_instance` | Launch a new GPU instance | +| `start_instance` | Launch a new GPU instance (auto-notifies if configured) | | `stop_instance` | Terminate a running instance | | `list_running_instances` | Show all running instances with status and connection details | | `check_availability` | Check if a specific GPU type is available | +### Auto-Notifications + +When notification environment variables are configured, the MCP server automatically sends notifications when instances become SSH-able. No additional flags needed—just set the `LAMBDA_NOTIFY_*` env vars and launch instances as usual. + ### Claude Code Setup ```bash @@ -157,7 +202,7 @@ Then restart Claude Code. Once configured, you can ask Claude things like: -- "What GPUs are currently available on Lambda Labs?" +- "What GPUs are currently available on Lambda?" - "Launch an H100 instance with my ssh key 'macbook'" - "Show me my running instances" - "Check if any A100s are available" diff --git a/npm/lambda-mcp/package.json b/npm/lambda-mcp/package.json index 6671d09..8ee9762 100644 --- a/npm/lambda-mcp/package.json +++ b/npm/lambda-mcp/package.json @@ -1,7 +1,7 @@ { "name": "@strand-ai/lambda-mcp", "version": "0.0.0", - "description": "MCP server for Lambda Labs cloud GPU instances", + "description": "Unofficial MCP server for Lambda cloud GPU instances", "bin": "lib/index.js", "main": "lib/index.js", "type": "commonjs", diff --git a/scripts/update-homebrew.sh b/scripts/update-homebrew.sh index 545f80c..0c83365 100755 --- a/scripts/update-homebrew.sh +++ b/scripts/update-homebrew.sh @@ -46,7 +46,7 @@ echo "Updating formula at $FORMULA_PATH..." cat > "$FORMULA_PATH" << EOF class LambdaCli < Formula - desc "CLI tool for Lambda Labs cloud GPU API" + desc "Unofficial CLI tool for Lambda cloud GPU API" homepage "https://github.com/Strand-AI/lambda-cli" url "https://github.com/Strand-AI/lambda-cli/archive/refs/tags/v${VERSION}.tar.gz" sha256 "${SHA256}" diff --git a/src/api.rs b/src/api.rs index 86e7fcd..4760438 100644 --- a/src/api.rs +++ b/src/api.rs @@ -135,7 +135,7 @@ enum ApiKeySource { Command(String), } -/// Lambda Labs API client +/// Lambda API client pub struct LambdaClient { client: Client, api_key_source: ApiKeySource, @@ -244,7 +244,7 @@ impl LambdaClient { .header(AUTHORIZATION, format!("Bearer {}", api_key)) .send() .await - .context("Failed to connect to Lambda Labs API")?; + .context("Failed to connect to Lambda API")?; if response.status().is_success() { Ok(()) diff --git a/src/bin/lambda-mcp.rs b/src/bin/lambda-mcp.rs index c0456a7..8125179 100644 --- a/src/bin/lambda-mcp.rs +++ b/src/bin/lambda-mcp.rs @@ -1,5 +1,6 @@ use anyhow::Result; use lambda_cli::api::{Filesystem, Instance, InstanceTypeData, LambdaClient}; +use lambda_cli::notify::{InstanceReadyMessage, Notifier, NotifyConfig}; use rmcp::handler::server::tool::ToolRouter; use rmcp::handler::server::wrapper::Parameters; use rmcp::model::{CallToolResult, Content, ServerCapabilities, ServerInfo}; @@ -7,11 +8,13 @@ use rmcp::schemars::JsonSchema; use rmcp::serde::Deserialize; use rmcp::{tool, tool_handler, tool_router, ErrorData as McpError, ServerHandler, ServiceExt}; use std::sync::Arc; +use std::time::Duration; -/// Lambda Labs MCP Server +/// Lambda MCP Server #[derive(Clone)] struct LambdaService { client: Arc, + notify_config: Option, #[allow(dead_code)] tool_router: ToolRouter, } @@ -20,8 +23,21 @@ impl LambdaService { fn new(lazy: bool) -> Result { dotenv::dotenv().ok(); let client = LambdaClient::from_env_with_options(lazy)?; + let notify_config = NotifyConfig::from_env(); + + // Debug: log notification config status + if let Some(ref config) = notify_config { + eprintln!( + "[lambda-mcp] Notifications configured for: {}", + config.configured_channels().join(", ") + ); + } else { + eprintln!("[lambda-mcp] No notification channels configured"); + } + Ok(Self { client: Arc::new(client), + notify_config, tool_router: Self::tool_router(), }) } @@ -166,7 +182,9 @@ impl LambdaService { )])) } - #[tool(description = "Launch a new GPU instance. Returns instance ID and connection details. Optionally attach a filesystem (must be in the same region).")] + #[tool( + description = "Launch a new GPU instance. Returns instance ID and connection details. Optionally attach a filesystem (must be in the same region). If notification env vars are configured, will auto-notify when instance is SSH-able." + )] async fn start_instance( &self, Parameters(params): Parameters, @@ -189,9 +207,36 @@ impl LambdaService { .map(|f| format!("\nFilesystem: {} (mounted at /lambda/nfs/{})", f, f)) .unwrap_or_default(); + // Spawn background task to notify when instance is ready + let notify_status = if let Some(ref config) = self.notify_config { + let channels = config.configured_channels().join(", "); + let client = Arc::clone(&self.client); + let notifier = Notifier::new(config.clone()); + let instance_id = result.instance_id.clone(); + let instance_name = params.name.clone(); + let gpu_type = params.gpu.clone(); + let region = result.region.clone(); + + tokio::spawn(async move { + poll_and_notify( + client, + notifier, + instance_id, + instance_name, + gpu_type, + region, + ) + .await; + }); + + format!("\n\nNotifications enabled for: {}. You will be notified when the instance is SSH-able.", channels) + } else { + String::new() + }; + Ok(CallToolResult::success(vec![Content::text(format!( - "Instance launched successfully!\n\nInstance ID: {}\nRegion: {}{}\n\nNote: Instance may take a few minutes to become active. Use 'list_running_instances' to check status.", - result.instance_id, result.region, fs_info + "Instance launched successfully!\n\nInstance ID: {}\nRegion: {}{}\n\nNote: Instance may take a few minutes to become active. Use 'list_running_instances' to check status.{}", + result.instance_id, result.region, fs_info, notify_status ))])) } @@ -311,7 +356,7 @@ impl ServerHandler for LambdaService { fn get_info(&self) -> ServerInfo { ServerInfo { instructions: Some( - "Lambda Labs GPU instance management. Use these tools to list available GPU types, \ + "Lambda GPU instance management. Use these tools to list available GPU types, \ launch instances, check running instances, terminate instances, and manage \ persistent filesystems. Filesystems can be attached to instances at launch time \ for persistent storage. Requires LAMBDA_API_KEY environment variable to be set." @@ -323,6 +368,76 @@ impl ServerHandler for LambdaService { } } +/// Background task to poll for instance readiness and send notifications +async fn poll_and_notify( + client: Arc, + notifier: Notifier, + instance_id: String, + instance_name: Option, + gpu_type: String, + region: String, +) { + let max_wait = Duration::from_secs(600); // 10 minutes max + let poll_interval = Duration::from_secs(10); + let start = std::time::Instant::now(); + + loop { + if start.elapsed() > max_wait { + eprintln!( + "[notify] Timeout waiting for instance {} to become active", + instance_id + ); + return; + } + + tokio::time::sleep(poll_interval).await; + + match client.get_instance(&instance_id).await { + Ok(instance) => { + let status = instance.status.as_deref().unwrap_or("unknown"); + + if status == "terminated" || status == "unhealthy" { + eprintln!( + "[notify] Instance {} entered {} state, stopping notifications", + instance_id, status + ); + return; + } + + // Notify when IP is available (don't wait for "active" status) + if let Some(ip) = instance.ip { + let msg = InstanceReadyMessage { + instance_id: instance_id.clone(), + instance_name, + ip, + gpu_type, + region, + }; + + let results = notifier.send_all(&msg).await; + for (channel, result) in results { + match result { + Ok(()) => eprintln!( + "[notify] {} notification sent for {}", + channel, instance_id + ), + Err(e) => { + eprintln!("[notify] {} notification failed: {}", channel, e) + } + } + } + return; + } + // No IP yet, continue polling + } + Err(e) => { + eprintln!("[notify] Error checking instance {}: {}", instance_id, e); + // Continue polling on transient errors + } + } + } +} + #[tokio::main] async fn main() -> Result<()> { // Parse command line arguments diff --git a/src/lib.rs b/src/lib.rs index 34ba9a8..7d9d069 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ pub mod api; +pub mod notify; pub use api::{Filesystem, Instance, InstanceTypeData, LambdaClient, LambdaError, LaunchResult}; +pub use notify::{InstanceReadyMessage, Notifier, NotifyConfig}; diff --git a/src/main.rs b/src/main.rs index 1465e5e..c0e6c7f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,16 +8,17 @@ use crossterm::{ terminal::{Clear, ClearType}, }; use lambda_cli::api::{LambdaClient, LambdaError}; +use lambda_cli::notify::{InstanceReadyMessage, Notifier, NotifyConfig}; use prettytable::{row, Table}; use std::io::{stdout, Write}; use std::time::{Duration, Instant}; use tokio::runtime::Runtime; -/// A command-line tool for Lambda Labs cloud GPU API +/// A command-line tool for Lambda cloud GPU API #[derive(Parser)] #[command(name = "lambda")] #[command(version = "0.2.0")] -#[command(about = "A command-line tool for Lambda Labs cloud GPU API", long_about = None)] +#[command(about = "A command-line tool for Lambda cloud GPU API", long_about = None)] struct Cli { #[command(subcommand)] command: Option, @@ -44,6 +45,9 @@ enum Commands { /// Filesystem name to attach (must be in same region) #[arg(short, long)] filesystem: Option, + /// Disable notifications even if LAMBDA_NOTIFY_* env vars are set + #[arg(long)] + no_notify: bool, }, /// Stop a specified GPU instance Stop { @@ -70,6 +74,9 @@ enum Commands { /// Filesystem name to attach when launched (must be in same region) #[arg(short, long)] filesystem: Option, + /// Disable notifications even if LAMBDA_NOTIFY_* env vars are set + #[arg(long)] + no_notify: bool, }, /// List all filesystems (persistent storage) Filesystems, @@ -112,6 +119,7 @@ fn run() -> Result<()> { name, region, filesystem, + no_notify, }) => start_instance( &rt, &client, @@ -120,6 +128,7 @@ fn run() -> Result<()> { name.as_deref(), region.as_deref(), filesystem.as_deref(), + *no_notify, ), Some(Commands::Stop { instance_id }) => stop_instance(&rt, &client, instance_id), Some(Commands::Running) => list_running_instances(&rt, &client), @@ -129,6 +138,7 @@ fn run() -> Result<()> { interval, name, filesystem, + no_notify, }) => find_and_start_instance( &rt, &client, @@ -137,6 +147,7 @@ fn run() -> Result<()> { *interval, name.as_deref(), filesystem.as_deref(), + *no_notify, ), Some(Commands::Filesystems) => list_filesystems(&rt, &client), Some(Commands::CreateFilesystem { name, region }) => { @@ -197,6 +208,7 @@ fn list_instances(rt: &Runtime, client: &LambdaClient) -> Result<()> { Ok(()) } +#[allow(clippy::too_many_arguments)] fn start_instance( rt: &Runtime, client: &LambdaClient, @@ -205,7 +217,22 @@ fn start_instance( name: Option<&str>, region: Option<&str>, filesystem: Option<&str>, + no_notify: bool, ) -> Result<()> { + // Auto-enable notifications if env vars are configured (unless --no-notify) + let notifier = if no_notify { + None + } else { + NotifyConfig::from_env().map(|config| { + println!( + "{} Notifications enabled for: {}", + "Info:".blue(), + config.configured_channels().join(", ") + ); + Notifier::new(config) + }) + }; + let fs_info = filesystem .map(|f| format!(" with filesystem '{}'", f.magenta())) .unwrap_or_default(); @@ -217,9 +244,8 @@ fn start_instance( fs_info ); - let result = rt.block_on(client.launch_instance_with_filesystem( - gpu, ssh, name, region, filesystem, - ))?; + let result = + rt.block_on(client.launch_instance_with_filesystem(gpu, ssh, name, region, filesystem))?; println!( "{} Instance {} launched in region {}", @@ -255,19 +281,40 @@ fn start_instance( ); stdout().flush().ok(); - if status == "active" { + // Ready when IP is available (don't wait for "active" status) + if let Some(ref ip) = instance.ip { println!(); - if let Some(ip) = instance.ip { - println!( - "{} Instance is active! SSH: {}", - "Ready!".green().bold(), - format!("ssh ubuntu@{}", ip).cyan() - ); - } else { - println!( - "{} Instance is active but IP not yet assigned", - "Ready!".green().bold() - ); + println!( + "{} Instance is ready! SSH: {}", + "Ready!".green().bold(), + format!("ssh ubuntu@{}", ip).cyan() + ); + + // Send notification if configured + if let Some(ref notifier) = notifier { + let msg = InstanceReadyMessage { + instance_id: result.instance_id.clone(), + instance_name: name.map(String::from), + ip: ip.clone(), + gpu_type: gpu.to_string(), + region: result.region.clone(), + }; + + println!("{} Sending notifications...", "Info:".blue()); + let results = rt.block_on(notifier.send_all(&msg)); + for (channel, result) in results { + match result { + Ok(()) => { + println!(" {} {} notification sent", "✓".green(), channel) + } + Err(e) => println!( + " {} {} notification failed: {}", + "✗".red(), + channel, + e + ), + } + } } break; } else if status == "terminated" || status == "unhealthy" { @@ -351,6 +398,7 @@ fn list_running_instances(rt: &Runtime, client: &LambdaClient) -> Result<()> { Ok(()) } +#[allow(clippy::too_many_arguments)] fn find_and_start_instance( rt: &Runtime, client: &LambdaClient, @@ -359,6 +407,7 @@ fn find_and_start_instance( interval: u64, name: Option<&str>, filesystem: Option<&str>, + no_notify: bool, ) -> Result<()> { if ssh.is_empty() { return Err(LambdaError::SshKeyRequired.into()); @@ -391,7 +440,7 @@ fn find_and_start_instance( regions.join(", ").blue() ); - return start_instance(rt, client, gpu, ssh, name, None, filesystem); + return start_instance(rt, client, gpu, ssh, name, None, filesystem, no_notify); } Ok(_) => { // No availability diff --git a/src/notify.rs b/src/notify.rs new file mode 100644 index 0000000..3d74db6 --- /dev/null +++ b/src/notify.rs @@ -0,0 +1,380 @@ +//! Notification module for Slack, Discord, and Telegram integrations. +//! +//! Configure via environment variables: +//! - `LAMBDA_NOTIFY_SLACK_WEBHOOK` - Slack incoming webhook URL +//! - `LAMBDA_NOTIFY_DISCORD_WEBHOOK` - Discord webhook URL +//! - `LAMBDA_NOTIFY_TELEGRAM_BOT_TOKEN` - Telegram bot token +//! - `LAMBDA_NOTIFY_TELEGRAM_CHAT_ID` - Telegram chat ID + +use anyhow::Result; +use reqwest::Client; +use serde_json::json; +use std::time::Duration; + +/// Message payload for instance ready notifications +#[derive(Debug, Clone)] +pub struct InstanceReadyMessage { + pub instance_id: String, + pub instance_name: Option, + pub ip: String, + pub gpu_type: String, + pub region: String, +} + +impl InstanceReadyMessage { + pub fn ssh_command(&self) -> String { + format!("ssh ubuntu@{}", self.ip) + } + + pub fn display_name(&self) -> &str { + self.instance_name.as_deref().unwrap_or(&self.instance_id) + } +} + +/// Slack webhook configuration +#[derive(Debug, Clone)] +pub struct SlackConfig { + pub webhook_url: String, +} + +/// Discord webhook configuration +#[derive(Debug, Clone)] +pub struct DiscordConfig { + pub webhook_url: String, +} + +/// Telegram bot configuration +#[derive(Debug, Clone)] +pub struct TelegramConfig { + pub bot_token: String, + pub chat_id: String, +} + +/// Combined notification configuration +#[derive(Debug, Clone, Default)] +pub struct NotifyConfig { + pub slack: Option, + pub discord: Option, + pub telegram: Option, +} + +impl NotifyConfig { + /// Load notification configuration from environment variables + pub fn from_env() -> Option { + let slack = std::env::var("LAMBDA_NOTIFY_SLACK_WEBHOOK") + .ok() + .filter(|s| !s.is_empty()) + .map(|webhook_url| SlackConfig { webhook_url }); + + let discord = std::env::var("LAMBDA_NOTIFY_DISCORD_WEBHOOK") + .ok() + .filter(|s| !s.is_empty()) + .map(|webhook_url| DiscordConfig { webhook_url }); + + let telegram = match ( + std::env::var("LAMBDA_NOTIFY_TELEGRAM_BOT_TOKEN"), + std::env::var("LAMBDA_NOTIFY_TELEGRAM_CHAT_ID"), + ) { + (Ok(bot_token), Ok(chat_id)) if !bot_token.is_empty() && !chat_id.is_empty() => { + Some(TelegramConfig { bot_token, chat_id }) + } + _ => None, + }; + + if slack.is_some() || discord.is_some() || telegram.is_some() { + Some(Self { + slack, + discord, + telegram, + }) + } else { + None + } + } + + /// Check if any notification channel is configured + pub fn is_configured(&self) -> bool { + self.slack.is_some() || self.discord.is_some() || self.telegram.is_some() + } + + /// Get a list of configured notification channels (for display) + pub fn configured_channels(&self) -> Vec<&'static str> { + let mut channels = Vec::new(); + if self.slack.is_some() { + channels.push("Slack"); + } + if self.discord.is_some() { + channels.push("Discord"); + } + if self.telegram.is_some() { + channels.push("Telegram"); + } + channels + } +} + +/// Notifier for sending messages to configured channels +pub struct Notifier { + client: Client, + config: NotifyConfig, +} + +impl Notifier { + /// Create a new notifier with the given configuration + pub fn new(config: NotifyConfig) -> Self { + let client = Client::builder() + .timeout(Duration::from_secs(10)) + .build() + .expect("Failed to create HTTP client"); + + Self { client, config } + } + + /// Create a notifier from environment variables, if configured + pub fn from_env() -> Option { + NotifyConfig::from_env().map(Self::new) + } + + /// Send notification to all configured channels + /// + /// Returns a list of (channel_name, result) for each attempt + pub async fn send_all(&self, msg: &InstanceReadyMessage) -> Vec<(&'static str, Result<()>)> { + let mut results = Vec::new(); + + if let Some(ref slack) = self.config.slack { + let result = self.send_slack(slack, msg).await; + results.push(("Slack", result)); + } + + if let Some(ref discord) = self.config.discord { + let result = self.send_discord(discord, msg).await; + results.push(("Discord", result)); + } + + if let Some(ref telegram) = self.config.telegram { + let result = self.send_telegram(telegram, msg).await; + results.push(("Telegram", result)); + } + + results + } + + /// Send notification to Slack + async fn send_slack(&self, config: &SlackConfig, msg: &InstanceReadyMessage) -> Result<()> { + let payload = json!({ + "blocks": [ + { + "type": "header", + "text": { + "type": "plain_text", + "text": "GPU Instance Ready!", + "emoji": true + } + }, + { + "type": "section", + "fields": [ + { + "type": "mrkdwn", + "text": format!("*Name:*\n{}", msg.display_name()) + }, + { + "type": "mrkdwn", + "text": format!("*GPU:*\n{}", msg.gpu_type) + }, + { + "type": "mrkdwn", + "text": format!("*Region:*\n{}", msg.region) + }, + { + "type": "mrkdwn", + "text": format!("*IP:*\n{}", msg.ip) + } + ] + }, + { + "type": "section", + "text": { + "type": "mrkdwn", + "text": format!("*SSH Command:*\n```{}```", msg.ssh_command()) + } + } + ] + }); + + let response = self + .client + .post(&config.webhook_url) + .json(&payload) + .send() + .await?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("Slack webhook failed ({}): {}", status, body); + } + + Ok(()) + } + + /// Send notification to Discord + async fn send_discord(&self, config: &DiscordConfig, msg: &InstanceReadyMessage) -> Result<()> { + let payload = json!({ + "embeds": [{ + "title": "GPU Instance Ready!", + "color": 5763719, // Green color + "fields": [ + { + "name": "Name", + "value": msg.display_name(), + "inline": true + }, + { + "name": "GPU", + "value": msg.gpu_type, + "inline": true + }, + { + "name": "Region", + "value": msg.region, + "inline": true + }, + { + "name": "IP Address", + "value": msg.ip, + "inline": true + }, + { + "name": "SSH Command", + "value": format!("```{}```", msg.ssh_command()), + "inline": false + } + ] + }] + }); + + let response = self + .client + .post(&config.webhook_url) + .json(&payload) + .send() + .await?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("Discord webhook failed ({}): {}", status, body); + } + + Ok(()) + } + + /// Send notification to Telegram + async fn send_telegram( + &self, + config: &TelegramConfig, + msg: &InstanceReadyMessage, + ) -> Result<()> { + let url = format!( + "https://api.telegram.org/bot{}/sendMessage", + config.bot_token + ); + + let text = format!( + "*GPU Instance Ready\\!*\n\n\ + *Name:* `{}`\n\ + *GPU:* {}\n\ + *Region:* {}\n\ + *IP:* `{}`\n\n\ + *SSH Command:*\n```\n{}\n```", + escape_telegram_markdown(msg.display_name()), + escape_telegram_markdown(&msg.gpu_type), + escape_telegram_markdown(&msg.region), + msg.ip, + msg.ssh_command() + ); + + let payload = json!({ + "chat_id": config.chat_id, + "parse_mode": "MarkdownV2", + "text": text + }); + + let response = self.client.post(&url).json(&payload).send().await?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("Telegram API failed ({}): {}", status, body); + } + + Ok(()) + } +} + +/// Escape special characters for Telegram MarkdownV2 +fn escape_telegram_markdown(text: &str) -> String { + let special_chars = [ + '_', '*', '[', ']', '(', ')', '~', '`', '>', '#', '+', '-', '=', '|', '{', '}', '.', '!', + ]; + let mut result = String::with_capacity(text.len() * 2); + for c in text.chars() { + if special_chars.contains(&c) { + result.push('\\'); + } + result.push(c); + } + result +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_escape_telegram_markdown() { + assert_eq!(escape_telegram_markdown("hello"), "hello"); + assert_eq!(escape_telegram_markdown("hello_world"), "hello\\_world"); + assert_eq!(escape_telegram_markdown("gpu_1x_a100"), "gpu\\_1x\\_a100"); + } + + #[test] + fn test_instance_ready_message() { + let msg = InstanceReadyMessage { + instance_id: "abc123".to_string(), + instance_name: Some("my-gpu".to_string()), + ip: "1.2.3.4".to_string(), + gpu_type: "gpu_1x_a100".to_string(), + region: "us-east-1".to_string(), + }; + + assert_eq!(msg.ssh_command(), "ssh ubuntu@1.2.3.4"); + assert_eq!(msg.display_name(), "my-gpu"); + + let msg_no_name = InstanceReadyMessage { + instance_id: "abc123".to_string(), + instance_name: None, + ip: "1.2.3.4".to_string(), + gpu_type: "gpu_1x_a100".to_string(), + region: "us-east-1".to_string(), + }; + + assert_eq!(msg_no_name.display_name(), "abc123"); + } + + #[test] + fn test_notify_config_is_configured() { + let empty = NotifyConfig::default(); + assert!(!empty.is_configured()); + + let with_slack = NotifyConfig { + slack: Some(SlackConfig { + webhook_url: "https://hooks.slack.com/test".to_string(), + }), + ..Default::default() + }; + assert!(with_slack.is_configured()); + assert_eq!(with_slack.configured_channels(), vec!["Slack"]); + } +}