diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ea6ef7a..3050684 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -59,16 +59,13 @@ jobs: include: - os: ubuntu-latest target: x86_64-unknown-linux-gnu - artifact_name: lambda_cli - asset_name: lambda_cli-x86_64-unknown-linux-gnu + asset_name: lambda-cli-x86_64-unknown-linux-gnu - os: macos-latest target: x86_64-apple-darwin - artifact_name: lambda_cli - asset_name: lambda_cli-x86_64-apple-darwin + asset_name: lambda-cli-x86_64-apple-darwin - os: macos-latest target: aarch64-apple-darwin - artifact_name: lambda_cli - asset_name: lambda_cli-aarch64-apple-darwin + asset_name: lambda-cli-aarch64-apple-darwin steps: - uses: actions/checkout@v4 @@ -94,7 +91,7 @@ jobs: - name: Create tarball run: | cd target/${{ matrix.target }}/release - tar -czvf ../../../${{ matrix.asset_name }}.tar.gz ${{ matrix.artifact_name }} + tar -czvf ../../../${{ matrix.asset_name }}.tar.gz lambda lambda-mcp cd ../../.. - name: Upload artifact diff --git a/.github/workflows/release-please.yml b/.github/workflows/release-please.yml index 6054a0b..fdd74f0 100644 --- a/.github/workflows/release-please.yml +++ b/.github/workflows/release-please.yml @@ -68,16 +68,24 @@ jobs: include: - os: ubuntu-latest target: x86_64-unknown-linux-gnu - artifact_name: lambda_cli - asset_name: lambda_cli-x86_64-unknown-linux-gnu + asset_name: lambda-cli-x86_64-unknown-linux-gnu + npm_os: linux + npm_arch: x64 - os: macos-latest target: x86_64-apple-darwin - artifact_name: lambda_cli - asset_name: lambda_cli-x86_64-apple-darwin + asset_name: lambda-cli-x86_64-apple-darwin + npm_os: darwin + npm_arch: x64 - os: macos-latest target: aarch64-apple-darwin - artifact_name: lambda_cli - asset_name: lambda_cli-aarch64-apple-darwin + asset_name: lambda-cli-aarch64-apple-darwin + npm_os: darwin + npm_arch: arm64 + - os: windows-latest + target: x86_64-pc-windows-msvc + asset_name: lambda-cli-x86_64-pc-windows-msvc + npm_os: win32 + npm_arch: x64 steps: - uses: actions/checkout@v4 with: @@ -91,16 +99,54 @@ jobs: - name: Build release run: cargo build --release --target ${{ matrix.target }} - - name: Create tarball + - name: Create tarball (Unix) + if: matrix.os != 'windows-latest' run: | cd target/${{ matrix.target }}/release - tar -czvf ../../../${{ matrix.asset_name }}.tar.gz ${{ matrix.artifact_name }} + tar -czvf ../../../${{ matrix.asset_name }}.tar.gz lambda lambda-mcp + + - name: Create tarball (Windows) + if: matrix.os == 'windows-latest' + run: | + cd target/${{ matrix.target }}/release + 7z a ../../../${{ matrix.asset_name }}.zip lambda.exe lambda-mcp.exe + shell: bash - name: Upload artifact uses: actions/upload-artifact@v4 with: name: ${{ matrix.asset_name }} - path: ${{ matrix.asset_name }}.tar.gz + path: ${{ matrix.asset_name }}.${{ matrix.os == 'windows-latest' && 'zip' || 'tar.gz' }} + + - name: Create npm platform package + run: | + mkdir -p npm-pkg/bin + if [ "${{ matrix.os }}" = "windows-latest" ]; then + cp target/${{ matrix.target }}/release/lambda-mcp.exe npm-pkg/bin/ + else + cp target/${{ matrix.target }}/release/lambda-mcp npm-pkg/bin/ + fi + cat > npm-pkg/package.json << EOF + { + "name": "@strand-ai/lambda-mcp-${{ matrix.npm_os }}-${{ matrix.npm_arch }}", + "version": "${{ needs.release.outputs.version }}", + "description": "Platform-specific binary for lambda-mcp (${{ matrix.npm_os }}-${{ matrix.npm_arch }})", + "os": ["${{ matrix.npm_os }}"], + "cpu": ["${{ matrix.npm_arch }}"], + "repository": { + "type": "git", + "url": "https://github.com/Strand-AI/lambda-cli" + }, + "license": "MIT" + } + EOF + shell: bash + + - name: Upload npm package artifact + uses: actions/upload-artifact@v4 + with: + name: npm-lambda-mcp-${{ matrix.npm_os }}-${{ matrix.npm_arch }} + path: npm-pkg/ publish: needs: [release, build] @@ -115,7 +161,9 @@ jobs: uses: softprops/action-gh-release@v2 with: tag_name: v${{ needs.release.outputs.version }} - files: artifacts/**/*.tar.gz + files: | + artifacts/**/*.tar.gz + artifacts/**/*.zip generate_release_notes: true update-homebrew: @@ -140,7 +188,7 @@ jobs: run: | cat > homebrew-tap/Formula/lambda-cli.rb << 'EOF' class LambdaCli < Formula - desc "CLI tool for Lambda Labs cloud GPU API" + desc "CLI and MCP server for Lambda Labs 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 }}" @@ -150,11 +198,13 @@ jobs: depends_on "rust" => :build def install - system "cargo", "install", *std_cargo_args + system "cargo", "build", "--release" + bin.install "target/release/lambda" + bin.install "target/release/lambda-mcp" end test do - assert_match "lambda ${{ needs.release.outputs.version }}", shell_output("#{bin}/lambda_cli --version") + assert_match "lambda", shell_output("#{bin}/lambda --version") end end EOF @@ -167,3 +217,49 @@ jobs: git add Formula/lambda-cli.rb git commit -m "lambda-cli: update to ${{ needs.release.outputs.version }}" git push + + publish-npm: + needs: [release, build] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + ref: v${{ needs.release.outputs.version }} + + - uses: actions/setup-node@v4 + with: + node-version: '20' + registry-url: 'https://registry.npmjs.org' + + - name: Download npm package artifacts + uses: actions/download-artifact@v4 + with: + pattern: npm-lambda-mcp-* + path: npm-packages + + - name: Publish platform packages + env: + NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }} + run: | + for pkg_dir in npm-packages/npm-lambda-mcp-*/; do + echo "Publishing $pkg_dir" + cd "$pkg_dir" + npm publish --access public + cd - + done + + - name: Prepare and publish main package + env: + NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }} + run: | + cd npm/lambda-mcp + # Update version in package.json + sed -i 's/"version": "0.0.0"/"version": "${{ needs.release.outputs.version }}"/' package.json + # Update optional dependency versions + sed -i 's/"@strand-ai\/lambda-mcp-linux-x64": "0.0.0"/"@strand-ai\/lambda-mcp-linux-x64": "${{ needs.release.outputs.version }}"/' package.json + sed -i 's/"@strand-ai\/lambda-mcp-darwin-x64": "0.0.0"/"@strand-ai\/lambda-mcp-darwin-x64": "${{ needs.release.outputs.version }}"/' package.json + sed -i 's/"@strand-ai\/lambda-mcp-darwin-arm64": "0.0.0"/"@strand-ai\/lambda-mcp-darwin-arm64": "${{ needs.release.outputs.version }}"/' package.json + sed -i 's/"@strand-ai\/lambda-mcp-win32-x64": "0.0.0"/"@strand-ai\/lambda-mcp-win32-x64": "${{ needs.release.outputs.version }}"/' package.json + npm install + npm run build + npm publish --access public diff --git a/Cargo.lock b/Cargo.lock index 22e870f..415a8fd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,12 +17,6 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" -[[package]] -name = "android-tzdata" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" - [[package]] name = "android_system_properties" version = "0.1.5" @@ -156,16 +150,16 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" [[package]] name = "chrono" -version = "0.4.38" +version = "0.4.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a21f936df1771bf62b77f047b726c4625ff2e8aa607c01ec06e5a05bd8463401" +checksum = "145052bdd345b87320e369255277e3fb5152762ad123a901ef5c262dd38fe8d2" dependencies = [ - "android-tzdata", "iana-time-zone", "js-sys", "num-traits", + "serde", "wasm-bindgen", - "windows-targets 0.52.5", + "windows-link 0.2.1", ] [[package]] @@ -286,6 +280,41 @@ dependencies = [ "memchr", ] +[[package]] +name = "darling" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" +dependencies = [ + "darling_core", + "quote", + "syn", +] + [[package]] name = "dirs-next" version = "2.0.0" @@ -313,6 +342,12 @@ version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f" +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "encode_unicode" version = "1.0.0" @@ -386,11 +421,26 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" +checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" dependencies = [ "futures-core", "futures-sink", @@ -398,36 +448,60 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" +checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" + +[[package]] +name = "futures-executor" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] [[package]] name = "futures-io" -version = "0.3.30" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" + +[[package]] +name = "futures-macro" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" +checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] [[package]] name = "futures-sink" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fb8e00e87438d937621c1c6269e53f536c14d3fbd6a042bb24879e57d474fb5" +checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" [[package]] name = "futures-task" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004" +checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" [[package]] name = "futures-util" -version = "0.3.30" +version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" +checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" dependencies = [ + "futures-channel", "futures-core", "futures-io", + "futures-macro", "futures-sink", "futures-task", "memchr", @@ -490,6 +564,12 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "http" version = "1.4.0" @@ -630,6 +710,12 @@ dependencies = [ "cc", ] +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "idna" version = "0.5.0" @@ -672,7 +758,7 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f23ff5ef2b80d608d61efee834934d862cd92461afc0560dedf493e4c033738b" dependencies = [ - "hermit-abi", + "hermit-abi 0.3.9", "libc", "windows-sys 0.52.0", ] @@ -701,7 +787,7 @@ dependencies = [ [[package]] name = "lambda_cli" -version = "0.2.0" +version = "0.2.2" dependencies = [ "anyhow", "chrono", @@ -711,9 +797,12 @@ dependencies = [ "dotenv", "prettytable", "reqwest", + "rmcp", + "schemars", "serde", "serde_json", "thiserror 2.0.17", + "tokio", ] [[package]] @@ -831,6 +920,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi 0.5.2", + "libc", +] + [[package]] name = "object" version = "0.32.2" @@ -913,6 +1012,12 @@ dependencies = [ "windows-targets 0.52.5", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "percent-encoding" version = "2.3.1" @@ -989,6 +1094,26 @@ dependencies = [ "thiserror 1.0.60", ] +[[package]] +name = "ref-cast" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f354300ae66f76f1c85c5f84693f0ce81d747e2c3f21a45fef496d89c960bf7d" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "reqwest" version = "0.12.28" @@ -998,9 +1123,7 @@ dependencies = [ "base64", "bytes", "encoding_rs", - "futures-channel", "futures-core", - "futures-util", "h2", "http", "http-body", @@ -1045,6 +1168,40 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rmcp" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5947688160b56fb6c827e3c20a72c90392a1d7e9dec74749197aa1780ac42ca" +dependencies = [ + "base64", + "chrono", + "futures", + "paste", + "pin-project-lite", + "rmcp-macros", + "schemars", + "serde", + "serde_json", + "thiserror 2.0.17", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "rmcp-macros" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01263441d3f8635c628e33856c468b96ebbce1af2d3699ea712ca71432d4ee7a" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "serde_json", + "syn", +] + [[package]] name = "rustc-demangle" version = "0.1.24" @@ -1118,6 +1275,32 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "schemars" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54e910108742c57a770f492731f99be216a52fadd361b06c8fb59d74ccc267d2" +dependencies = [ + "chrono", + "dyn-clone", + "ref-cast", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4908ad288c5035a8eb12cfdf0d49270def0a268ee162b75eeee0f85d155a7c45" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -1149,18 +1332,39 @@ dependencies = [ [[package]] name = "serde" -version = "1.0.202" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "226b61a0d411b2ba5ff6d7f73a476ac4f8bb900373459cd00fab8512828ba395" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.202" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6048858004bcff69094cd972ed40a32500f153bd3be9f716b2eed2e8217c4838" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" dependencies = [ "proc-macro2", "quote", @@ -1169,13 +1373,15 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.117" +version = "1.0.149" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "455182ea6142b14f93f4bc5320a2b31c1f266b66a4a5c858b013302a5d8cbfc3" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" dependencies = [ "itoa", - "ryu", + "memchr", "serde", + "serde_core", + "zmij", ] [[package]] @@ -1402,11 +1608,26 @@ dependencies = [ "bytes", "libc", "mio 0.8.11", + "num_cpus", + "parking_lot", "pin-project-lite", + "signal-hook-registry", "socket2 0.5.7", + "tokio-macros", "windows-sys 0.48.0", ] +[[package]] +name = "tokio-macros" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b8a1e28f2deaa14e508979454cb3a223b10b938b45af148bc0986de36f1923b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tokio-native-tls" version = "0.3.1" @@ -1492,9 +1713,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" dependencies = [ "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tracing-core" version = "0.1.32" @@ -1947,3 +2180,9 @@ name = "zeroize" version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" + +[[package]] +name = "zmij" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd8f3f50b848df28f887acb68e41201b5aea6bc8a8dacc00fb40635ff9a72fea" diff --git a/Cargo.toml b/Cargo.toml index 33b9e8c..b24835d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,13 +2,21 @@ name = "lambda_cli" version = "0.2.2" edition = "2021" -description = "A command-line tool for Lambda Labs cloud GPU API" +description = "CLI and MCP server for Lambda Labs cloud GPU instances" license = "MIT" repository = "https://github.com/Strand-AI/lambda-cli" +[[bin]] +name = "lambda" +path = "src/main.rs" + +[[bin]] +name = "lambda-mcp" +path = "src/bin/lambda-mcp.rs" + [dependencies] clap = { version = "4.4", features = ["derive"] } -reqwest = { version = "0.12", features = ["json", "blocking"] } +reqwest = { version = "0.12", features = ["json"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" dotenv = "0.15" @@ -18,3 +26,6 @@ chrono = "0.4" crossterm = "0.28" anyhow = "1.0" thiserror = "2.0" +tokio = { version = "1", features = ["full"] } +rmcp = { version = "0.8", features = ["server", "macros", "transport-io"] } +schemars = "1" diff --git a/README.md b/README.md index 65c1bd3..d9249f3 100644 --- a/README.md +++ b/README.md @@ -3,63 +3,82 @@ [![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) -A fast CLI for managing [Lambda Labs](https://lambdalabs.com/) cloud GPU instances. +A fast CLI and MCP server for managing [Lambda Labs](https://lambdalabs.com/) cloud GPU instances. + +**Two ways to use it:** +- **CLI** (`lambda`) - Direct terminal commands for managing GPU instances +- **MCP Server** (`lambda-mcp`) - Let AI assistants like Claude manage your GPU infrastructure ## Installation -**Homebrew:** +### Homebrew (macOS/Linux) ```bash brew install strand-ai/tap/lambda-cli ``` -**From source:** +### From Source ```bash cargo install --git https://github.com/Strand-AI/lambda-cli ``` -## Setup +### Pre-built Binaries +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) and export it: +Get your API key from the [Lambda Labs dashboard](https://cloud.lambdalabs.com/api-keys). +### Option 1: Environment Variable ```bash export LAMBDA_API_KEY= ``` -## Commands +### Option 2: Command (1Password, etc.) +```bash +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. + +--- + +## CLI Usage + +### Commands | Command | Description | |---------|-------------| -| `list` | Show available GPU types with pricing and availability | -| `running` | Show your running instances | -| `start` | Launch a new instance | -| `stop` | Terminate an instance | -| `find` | Poll until a GPU type is available, then launch | +| `lambda list` | Show available GPU types with pricing and availability | +| `lambda running` | Show your running instances | +| `lambda start` | Launch a new instance | +| `lambda stop` | Terminate an instance | +| `lambda find` | Poll until a GPU type is available, then launch | -## Usage +### Examples **List available GPUs:** ```bash -lambda_cli list +lambda list ``` **Start an instance:** ```bash -lambda_cli start --gpu gpu_1x_a10 --ssh my-key --name "dev-box" +lambda start --gpu gpu_1x_a10 --ssh my-key --name "dev-box" ``` **Stop an instance:** ```bash -lambda_cli stop --instance-id +lambda stop --instance-id ``` **Wait for availability and auto-launch:** ```bash -lambda_cli find --gpu gpu_8x_h100 --ssh my-key --interval 30 +lambda find --gpu gpu_8x_h100 --ssh my-key --interval 30 ``` -## Options +### CLI Options -### start +#### start | Flag | Description | |------|-------------| | `-g, --gpu` | Instance type (required) | @@ -67,10 +86,110 @@ lambda_cli find --gpu gpu_8x_h100 --ssh my-key --interval 30 | `-n, --name` | Instance name | | `-r, --region` | Region (auto-selects if omitted) | -### find +#### find | Flag | Description | |------|-------------| | `-g, --gpu` | Instance type to wait for (required) | | `-s, --ssh` | SSH key name (required) | | `--interval` | Poll interval in seconds (default: 10) | | `-n, --name` | Instance name when launched | + +--- + +## 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. + +### Quick Start with npx + +The easiest way to use `lambda-mcp` is via npx—no installation required: + +```bash +npx @strand-ai/lambda-mcp +``` + +### Available Tools + +| Tool | Description | +|------|-------------| +| `list_gpu_types` | List all available GPU instance types with pricing, specs, and current availability | +| `start_instance` | Launch a new GPU instance | +| `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 | + +### Claude Code Setup + +Add to `~/.claude/claude_desktop_config.json`: + +**Using npx (no installation needed):** +```json +{ + "mcpServers": { + "lambda-labs": { + "command": "npx", + "args": ["-y", "@strand-ai/lambda-mcp"], + "env": { + "LAMBDA_API_KEY": "your-api-key-here" + } + } + } +} +``` + +**Using installed binary:** +```json +{ + "mcpServers": { + "lambda-labs": { + "command": "lambda-mcp", + "env": { + "LAMBDA_API_KEY": "your-api-key-here" + } + } + } +} +``` + +**With 1Password:** +```json +{ + "mcpServers": { + "lambda-labs": { + "command": "npx", + "args": ["-y", "@strand-ai/lambda-mcp"], + "env": { + "LAMBDA_API_KEY_COMMAND": "op read op://Personal/Lambda/api-key" + } + } + } +} +``` + +### Example Prompts + +Once configured, you can ask Claude things like: + +- "What GPUs are currently available on Lambda Labs?" +- "Launch an H100 instance with my ssh key 'macbook'" +- "Show me my running instances" +- "Check if any A100s are available" +- "Terminate instance i-abc123" + +--- + +## Development + +```bash +# Build +cargo build + +# Run tests +cargo test + +# Run CLI +cargo run --bin lambda -- list + +# Run MCP server +cargo run --bin lambda-mcp +``` diff --git a/npm/lambda-mcp/.gitignore b/npm/lambda-mcp/.gitignore new file mode 100644 index 0000000..88edb62 --- /dev/null +++ b/npm/lambda-mcp/.gitignore @@ -0,0 +1,2 @@ +node_modules/ +lib/ diff --git a/npm/lambda-mcp/package-lock.json b/npm/lambda-mcp/package-lock.json new file mode 100644 index 0000000..1b9c67d --- /dev/null +++ b/npm/lambda-mcp/package-lock.json @@ -0,0 +1,69 @@ +{ + "name": "@strand-ai/lambda-mcp", + "version": "0.0.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "@strand-ai/lambda-mcp", + "version": "0.0.0", + "license": "MIT", + "bin": { + "lambda-mcp": "lib/index.js" + }, + "devDependencies": { + "@types/node": "^20.0.0", + "typescript": "^5.0.0" + }, + "optionalDependencies": { + "@strand-ai/lambda-mcp-darwin-arm64": "0.0.0", + "@strand-ai/lambda-mcp-darwin-x64": "0.0.0", + "@strand-ai/lambda-mcp-linux-x64": "0.0.0", + "@strand-ai/lambda-mcp-win32-x64": "0.0.0" + } + }, + "node_modules/@strand-ai/lambda-mcp-darwin-arm64": { + "optional": true + }, + "node_modules/@strand-ai/lambda-mcp-darwin-x64": { + "optional": true + }, + "node_modules/@strand-ai/lambda-mcp-linux-x64": { + "optional": true + }, + "node_modules/@strand-ai/lambda-mcp-win32-x64": { + "optional": true + }, + "node_modules/@types/node": { + "version": "20.19.29", + "resolved": "https://registry.npmjs.org/@types/node/-/node-20.19.29.tgz", + "integrity": "sha512-YrT9ArrGaHForBaCNwFjoqJWmn8G1Pr7+BH/vwyLHciA9qT/wSiuOhxGCT50JA5xLvFBd6PIiGkE3afxcPE1nw==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~6.21.0" + } + }, + "node_modules/typescript": { + "version": "5.9.3", + "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", + "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "tsc": "bin/tsc", + "tsserver": "bin/tsserver" + }, + "engines": { + "node": ">=14.17" + } + }, + "node_modules/undici-types": { + "version": "6.21.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", + "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", + "dev": true, + "license": "MIT" + } + } +} diff --git a/npm/lambda-mcp/package.json b/npm/lambda-mcp/package.json new file mode 100644 index 0000000..6671d09 --- /dev/null +++ b/npm/lambda-mcp/package.json @@ -0,0 +1,40 @@ +{ + "name": "@strand-ai/lambda-mcp", + "version": "0.0.0", + "description": "MCP server for Lambda Labs cloud GPU instances", + "bin": "lib/index.js", + "main": "lib/index.js", + "type": "commonjs", + "scripts": { + "build": "tsc", + "prepublishOnly": "npm run build" + }, + "keywords": [ + "mcp", + "model-context-protocol", + "lambda-labs", + "gpu", + "cloud", + "ai", + "claude" + ], + "repository": { + "type": "git", + "url": "https://github.com/Strand-AI/lambda-cli" + }, + "homepage": "https://github.com/Strand-AI/lambda-cli#readme", + "bugs": { + "url": "https://github.com/Strand-AI/lambda-cli/issues" + }, + "license": "MIT", + "devDependencies": { + "@types/node": "^20.0.0", + "typescript": "^5.0.0" + }, + "optionalDependencies": { + "@strand-ai/lambda-mcp-linux-x64": "0.0.0", + "@strand-ai/lambda-mcp-darwin-x64": "0.0.0", + "@strand-ai/lambda-mcp-darwin-arm64": "0.0.0", + "@strand-ai/lambda-mcp-win32-x64": "0.0.0" + } +} diff --git a/npm/lambda-mcp/src/index.ts b/npm/lambda-mcp/src/index.ts new file mode 100644 index 0000000..a23e855 --- /dev/null +++ b/npm/lambda-mcp/src/index.ts @@ -0,0 +1,39 @@ +#!/usr/bin/env node + +import { spawnSync } from "child_process"; + +/** + * Returns the executable path which is located inside `node_modules` + * The naming convention is lambda-mcp-${os}-${arch} + * If the platform is `win32` or `cygwin`, executable will include a `.exe` extension. + */ +function getExePath(): string { + const arch = process.arch; + let os = process.platform as string; + let extension = ""; + + if (["win32", "cygwin"].includes(process.platform)) { + os = "win32"; + extension = ".exe"; + } + + try { + return require.resolve(`@strand-ai/lambda-mcp-${os}-${arch}/bin/lambda-mcp${extension}`); + } catch (e) { + throw new Error( + `Couldn't find lambda-mcp binary for ${os}-${arch}. ` + + `Please report this issue at https://github.com/Strand-AI/lambda-cli/issues` + ); + } +} + +/** + * Runs lambda-mcp with the given arguments + */ +function run(): void { + const args = process.argv.slice(2); + const result = spawnSync(getExePath(), args, { stdio: "inherit" }); + process.exit(result.status ?? 0); +} + +run(); diff --git a/npm/lambda-mcp/tsconfig.json b/npm/lambda-mcp/tsconfig.json new file mode 100644 index 0000000..fc208b1 --- /dev/null +++ b/npm/lambda-mcp/tsconfig.json @@ -0,0 +1,16 @@ +{ + "compilerOptions": { + "target": "ES2020", + "module": "commonjs", + "lib": ["ES2020"], + "outDir": "./lib", + "rootDir": "./src", + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "declaration": true + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "lib"] +} diff --git a/npm/package.json.tmpl b/npm/package.json.tmpl new file mode 100644 index 0000000..f88928b --- /dev/null +++ b/npm/package.json.tmpl @@ -0,0 +1,12 @@ +{ + "name": "lambda-mcp-${node_os}-${node_arch}", + "version": "${node_version}", + "description": "Platform-specific binary for lambda-mcp (${node_os}-${node_arch})", + "os": ["${node_os}"], + "cpu": ["${node_arch}"], + "repository": { + "type": "git", + "url": "https://github.com/Strand-AI/lambda-cli" + }, + "license": "MIT" +} diff --git a/src/api.rs b/src/api.rs new file mode 100644 index 0000000..69a744e --- /dev/null +++ b/src/api.rs @@ -0,0 +1,502 @@ +use anyhow::{anyhow, Context, Result}; +use reqwest::header::AUTHORIZATION; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::time::Duration; +use thiserror::Error; + +pub const API_BASE_URL: &str = "https://cloud.lambdalabs.com/api/v1"; +const DEFAULT_TIMEOUT_SECS: u64 = 30; + +#[derive(Error, Debug)] +pub enum LambdaError { + #[error("API key not set. Set LAMBDA_API_KEY or LAMBDA_API_KEY_COMMAND environment variable")] + ApiKeyNotSet, + #[error("Failed to execute API key command: {0}")] + ApiKeyCommandFailed(String), + #[error("Instance type '{0}' not found")] + InstanceTypeNotFound(String), + #[error("No regions available for instance type '{0}'")] + NoRegionsAvailable(String), + #[error("No instance IDs returned from launch request")] + NoInstanceIds, + #[error("API request failed: {0}")] + ApiError(String), + #[error("SSH key is required for this operation")] + SshKeyRequired, +} + +#[derive(Deserialize, Debug)] +pub struct ApiResponse { + pub data: T, +} + +#[derive(Deserialize, Debug)] +pub struct ApiErrorResponse { + pub error: ApiErrorDetail, +} + +#[derive(Deserialize, Debug)] +pub struct ApiErrorDetail { + pub message: String, +} + +#[derive(Deserialize, Debug, Clone, Serialize)] +pub struct Instance { + pub id: Option, + pub name: Option, + pub status: Option, + pub ip: Option, + pub ssh_key_names: Option>, + pub instance_type: Option, + pub region: Option, +} + +#[derive(Deserialize, Debug, Clone, Serialize)] +pub struct InstanceTypeInfo { + pub name: Option, +} + +#[derive(Deserialize, Debug, Clone, Serialize)] +pub struct RegionInfo { + pub name: Option, +} + +#[derive(Deserialize, Debug)] +pub struct LaunchResponse { + pub instance_ids: Vec, +} + +#[derive(Deserialize, Debug, Clone, Serialize)] +pub struct InstanceTypeData { + pub name: String, + pub description: String, + pub price_cents_per_hour: i32, + pub vcpus: u32, + pub memory_gib: u32, + pub storage_gib: u32, + pub regions_available: Vec, +} + +#[derive(Deserialize, Debug, Clone)] +pub struct InstanceTypeResponse { + pub instance_type: InstanceType, + pub regions_with_capacity_available: Vec, +} + +#[derive(Deserialize, Debug, Clone)] +pub struct InstanceType { + pub description: String, + pub price_cents_per_hour: i32, + pub specs: InstanceSpecs, +} + +#[derive(Deserialize, Debug, Clone)] +pub struct InstanceSpecs { + pub vcpus: u32, + pub memory_gib: u32, + pub storage_gib: u32, +} + +#[derive(Deserialize, Debug, Clone)] +pub struct Region { + pub name: String, + #[allow(dead_code)] + pub description: String, +} + +/// Lambda Labs API client +pub struct LambdaClient { + client: Client, + api_key: String, +} + +impl LambdaClient { + pub fn new(api_key: String) -> 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 }) + } + + /// Create a client using environment variables for the API key. + /// + /// 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`) + pub fn from_env() -> Result { + let api_key = get_api_key_from_env()?; + Self::new(api_key) + } + + /// Validate the API key by making a test request + pub async fn validate_api_key(&self) -> Result<()> { + let url = format!("{}/instances", API_BASE_URL); + let response = self + .client + .get(&url) + .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .send() + .await + .context("Failed to connect to Lambda Labs API")?; + + if response.status().is_success() { + Ok(()) + } else { + let status = response.status(); + let error_msg = Self::parse_error_response(response).await; + Err(anyhow!( + "API key validation failed ({}): {}", + status, + error_msg + )) + } + } + + /// List all available instance types + pub async fn list_instance_types(&self) -> Result> { + let url = format!("{}/instance-types", API_BASE_URL); + let response = self + .client + .get(&url) + .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .send() + .await + .context("Failed to fetch instance types")?; + + if !response.status().is_success() { + let error_msg = Self::parse_error_response(response).await; + return Err(anyhow!("Failed to list instances: {}", error_msg)); + } + + let response: ApiResponse> = response + .json() + .await + .context("Failed to parse instance types response")?; + + let mut result: Vec = response + .data + .into_iter() + .map(|(name, data)| InstanceTypeData { + name, + description: data.instance_type.description, + price_cents_per_hour: data.instance_type.price_cents_per_hour, + vcpus: data.instance_type.specs.vcpus, + memory_gib: data.instance_type.specs.memory_gib, + storage_gib: data.instance_type.specs.storage_gib, + regions_available: data + .regions_with_capacity_available + .into_iter() + .map(|r| r.name) + .collect(), + }) + .collect(); + + result.sort_by(|a, b| a.name.cmp(&b.name)); + Ok(result) + } + + /// Get instance type details (for checking availability) + pub async fn get_instance_type(&self, gpu: &str) -> Result> { + let url = format!("{}/instance-types", API_BASE_URL); + let response = self + .client + .get(&url) + .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .send() + .await + .context("Failed to fetch instance types")?; + + if !response.status().is_success() { + let error_msg = Self::parse_error_response(response).await; + return Err(anyhow!("Failed to get instance types: {}", error_msg)); + } + + let response: ApiResponse> = response + .json() + .await + .context("Failed to parse instance types")?; + + Ok(response.data.get(gpu).cloned()) + } + + /// Launch a new instance + pub async fn launch_instance( + &self, + gpu: &str, + ssh_key: &str, + name: Option<&str>, + region: Option<&str>, + ) -> Result { + let instance_type_response = self + .get_instance_type(gpu) + .await? + .ok_or_else(|| LambdaError::InstanceTypeNotFound(gpu.to_string()))?; + + let region_name = if let Some(r) = region { + // Validate the specified region is available + if !instance_type_response + .regions_with_capacity_available + .iter() + .any(|reg| reg.name == r) + { + return Err(anyhow!( + "Region '{}' is not available for instance type '{}'. Available regions: {}", + r, + gpu, + instance_type_response + .regions_with_capacity_available + .iter() + .map(|reg| reg.name.as_str()) + .collect::>() + .join(", ") + )); + } + r.to_string() + } else { + // Auto-select first available region + instance_type_response + .regions_with_capacity_available + .first() + .ok_or_else(|| LambdaError::NoRegionsAvailable(gpu.to_string()))? + .name + .clone() + }; + + let url = format!("{}/instance-operations/launch", API_BASE_URL); + + let mut payload = serde_json::json!({ + "region_name": region_name, + "instance_type_name": gpu, + "ssh_key_names": [ssh_key], + "quantity": 1 + }); + + if let Some(instance_name) = name { + payload["name"] = serde_json::Value::String(instance_name.to_string()); + } + + let response = self + .client + .post(&url) + .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .json(&payload) + .send() + .await + .context("Failed to send launch request")?; + + if !response.status().is_success() { + let error_msg = Self::parse_error_response(response).await; + return Err(anyhow!("Failed to launch instance: {}", error_msg)); + } + + let parsed_response: ApiResponse = response + .json() + .await + .context("Failed to parse launch response")?; + + let instance_id = parsed_response + .data + .instance_ids + .first() + .ok_or(LambdaError::NoInstanceIds)? + .clone(); + + Ok(LaunchResult { + instance_id, + region: region_name, + }) + } + + /// Terminate an instance + pub async fn terminate_instance(&self, instance_id: &str) -> Result<()> { + let url = format!("{}/instance-operations/terminate", API_BASE_URL); + let payload = serde_json::json!({ + "instance_ids": [instance_id] + }); + + let response = self + .client + .post(&url) + .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .json(&payload) + .send() + .await + .context("Failed to send terminate request")?; + + if !response.status().is_success() { + let error_msg = Self::parse_error_response(response).await; + return Err(anyhow!("Failed to terminate instance: {}", error_msg)); + } + + Ok(()) + } + + /// List all running instances + pub async fn list_running_instances(&self) -> Result> { + let url = format!("{}/instances", API_BASE_URL); + let response = self + .client + .get(&url) + .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .send() + .await + .context("Failed to fetch running instances")?; + + if !response.status().is_success() { + let error_msg = Self::parse_error_response(response).await; + return Err(anyhow!("Failed to list running instances: {}", error_msg)); + } + + let response: ApiResponse> = response + .json() + .await + .context("Failed to parse running instances response")?; + + Ok(response.data) + } + + /// Get details for a specific instance + pub async fn get_instance(&self, instance_id: &str) -> Result { + let url = format!("{}/instances/{}", API_BASE_URL, instance_id); + let response = self + .client + .get(&url) + .header(AUTHORIZATION, format!("Bearer {}", self.api_key)) + .send() + .await + .context("Failed to fetch instance details")?; + + if !response.status().is_success() { + let status = response.status(); + let error_msg = Self::parse_error_response(response).await; + return Err(anyhow!( + "Failed to get instance details ({}): {}", + status, + error_msg + )); + } + + let response: ApiResponse = response + .json() + .await + .context("Failed to parse instance details")?; + + Ok(response.data) + } + + /// Check if a GPU type is available + pub async fn check_availability(&self, gpu: &str) -> Result> { + let instance_type = self + .get_instance_type(gpu) + .await? + .ok_or_else(|| LambdaError::InstanceTypeNotFound(gpu.to_string()))?; + + Ok(instance_type + .regions_with_capacity_available + .into_iter() + .map(|r| r.name) + .collect()) + } + + async fn parse_error_response(response: reqwest::Response) -> String { + response + .json::() + .await + .map(|e| e.error.message) + .unwrap_or_else(|_| "Unknown error".to_string()) + } +} + +#[derive(Debug, Clone)] +pub struct LaunchResult { + pub instance_id: String, + 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; + + let output = if cfg!(target_os = "windows") { + Command::new("cmd").args(["/C", command]).output() + } else { + Command::new("sh").args(["-c", command]).output() + }; + + match output { + Ok(output) => { + if output.status.success() { + let key = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if key.is_empty() { + Err(LambdaError::ApiKeyCommandFailed( + "Command returned empty output".to_string(), + ) + .into()) + } else { + Ok(key) + } + } else { + let stderr = String::from_utf8_lossy(&output.stderr); + Err( + LambdaError::ApiKeyCommandFailed(format!("Command failed: {}", stderr.trim())) + .into(), + ) + } + } + Err(e) => Err(LambdaError::ApiKeyCommandFailed(format!( + "Failed to execute command: {}", + e + )) + .into()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_lambda_error_messages() { + assert_eq!( + LambdaError::ApiKeyNotSet.to_string(), + "API key not set. Set LAMBDA_API_KEY or LAMBDA_API_KEY_COMMAND environment variable" + ); + assert_eq!( + LambdaError::InstanceTypeNotFound("gpu_1x_a100".to_string()).to_string(), + "Instance type 'gpu_1x_a100' not found" + ); + } + + #[test] + fn test_api_base_url() { + assert_eq!(API_BASE_URL, "https://cloud.lambdalabs.com/api/v1"); + } +} diff --git a/src/bin/lambda-mcp.rs b/src/bin/lambda-mcp.rs new file mode 100644 index 0000000..47a4135 --- /dev/null +++ b/src/bin/lambda-mcp.rs @@ -0,0 +1,240 @@ +use anyhow::Result; +use lambda_cli::api::{Instance, InstanceTypeData, LambdaClient}; +use rmcp::handler::server::tool::ToolRouter; +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 std::sync::Arc; + +/// Lambda Labs MCP Server +#[derive(Clone)] +struct LambdaService { + client: Arc, + #[allow(dead_code)] + tool_router: ToolRouter, +} + +impl LambdaService { + fn new() -> Result { + dotenv::dotenv().ok(); + let client = LambdaClient::from_env()?; + Ok(Self { + client: Arc::new(client), + tool_router: Self::tool_router(), + }) + } + + fn format_instance_types(types: &[InstanceTypeData]) -> String { + let mut output = String::from("Available GPU Instance Types:\n\n"); + for t in types { + let price = t.price_cents_per_hour as f64 / 100.0; + let availability = if t.regions_available.is_empty() { + "No availability".to_string() + } else { + t.regions_available.join(", ") + }; + output.push_str(&format!( + "• {} - {}\n Price: ${:.2}/hr | vCPUs: {} | RAM: {} GiB | Storage: {} GiB\n Regions: {}\n\n", + t.name, t.description, price, t.vcpus, t.memory_gib, t.storage_gib, availability + )); + } + output + } + + fn format_instances(instances: &[Instance]) -> String { + if instances.is_empty() { + return "No running instances.".to_string(); + } + let mut output = String::from("Running Instances:\n\n"); + for inst in instances { + let id = inst.id.as_deref().unwrap_or("N/A"); + let name = inst.name.as_deref().unwrap_or("-"); + let status = inst.status.as_deref().unwrap_or("unknown"); + let ip = inst.ip.as_deref().unwrap_or("N/A"); + let inst_type = inst + .instance_type + .as_ref() + .and_then(|t| t.name.as_deref()) + .unwrap_or("N/A"); + let region = inst + .region + .as_ref() + .and_then(|r| r.name.as_deref()) + .unwrap_or("N/A"); + let ssh_keys = inst + .ssh_key_names + .as_ref() + .map(|k| k.join(", ")) + .unwrap_or_else(|| "N/A".to_string()); + + output.push_str(&format!( + "• ID: {}\n Name: {} | Type: {} | Region: {}\n Status: {} | IP: {}\n SSH Keys: {}\n\n", + id, name, inst_type, region, status, ip, ssh_keys + )); + } + output + } +} + +// Tool parameter types +#[derive(Debug, Deserialize, JsonSchema)] +struct StartInstanceParams { + /// GPU instance type (e.g., gpu_1x_h100, gpu_1x_a100) + gpu: String, + /// SSH key name to use for the instance + ssh_key: String, + /// Optional name for the instance + name: Option, + /// Optional region to launch in (auto-selects if not specified) + region: Option, +} + +#[derive(Debug, Deserialize, JsonSchema)] +struct StopInstanceParams { + /// Instance ID to terminate + instance_id: String, +} + +#[derive(Debug, Deserialize, JsonSchema)] +struct CheckAvailabilityParams { + /// GPU instance type to check (e.g., gpu_1x_h100) + gpu: String, +} + +#[tool_router] +impl LambdaService { + #[tool( + description = "List all available GPU instance types with pricing, specs, and current availability" + )] + async fn list_gpu_types(&self) -> Result { + let types = self + .client + .list_instance_types() + .await + .map_err(|e| McpError::internal_error(e.to_string(), None))?; + + Ok(CallToolResult::success(vec![Content::text( + Self::format_instance_types(&types), + )])) + } + + #[tool(description = "Launch a new GPU instance. Returns instance ID and connection details.")] + async fn start_instance( + &self, + Parameters(params): Parameters, + ) -> Result { + let result = self + .client + .launch_instance( + ¶ms.gpu, + ¶ms.ssh_key, + params.name.as_deref(), + params.region.as_deref(), + ) + .await + .map_err(|e| McpError::internal_error(e.to_string(), None))?; + + 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 + ))])) + } + + #[tool(description = "Terminate a running GPU instance")] + async fn stop_instance( + &self, + Parameters(params): Parameters, + ) -> Result { + self.client + .terminate_instance(¶ms.instance_id) + .await + .map_err(|e| McpError::internal_error(e.to_string(), None))?; + + Ok(CallToolResult::success(vec![Content::text(format!( + "Instance {} terminated successfully.", + params.instance_id + ))])) + } + + #[tool( + description = "List all currently running GPU instances with their status and connection details" + )] + async fn list_running_instances(&self) -> Result { + let instances = self + .client + .list_running_instances() + .await + .map_err(|e| McpError::internal_error(e.to_string(), None))?; + + Ok(CallToolResult::success(vec![Content::text( + Self::format_instances(&instances), + )])) + } + + #[tool( + description = "Check if a specific GPU type is currently available. Returns list of regions with availability." + )] + async fn check_availability( + &self, + Parameters(params): Parameters, + ) -> Result { + let regions = self + .client + .check_availability(¶ms.gpu) + .await + .map_err(|e| McpError::internal_error(e.to_string(), None))?; + + let message = if regions.is_empty() { + format!( + "GPU type '{}' is currently NOT available in any region.", + params.gpu + ) + } else { + format!( + "GPU type '{}' is available in: {}", + params.gpu, + regions.join(", ") + ) + }; + + Ok(CallToolResult::success(vec![Content::text(message)])) + } +} + +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, \ + launch instances, check running instances, and terminate instances. \ + Requires LAMBDA_API_KEY environment variable to be set." + .to_string(), + ), + capabilities: ServerCapabilities::builder().enable_tools().build(), + ..Default::default() + } + } +} + +#[tokio::main] +async fn main() -> Result<()> { + // Initialize the service + let service = match LambdaService::new() { + Ok(s) => s, + Err(e) => { + eprintln!("Failed to initialize Lambda service: {}", e); + eprintln!("Make sure LAMBDA_API_KEY environment variable is set."); + std::process::exit(1); + } + }; + + // Run the MCP server over stdio + let server = service.serve(rmcp::transport::io::stdio()).await?; + + // Wait for the server to finish + server.waiting().await?; + + Ok(()) +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..bfb1a28 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,3 @@ +pub mod api; + +pub use api::{Instance, InstanceTypeData, LambdaClient, LambdaError, LaunchResult}; diff --git a/src/main.rs b/src/main.rs index d5321c9..56f91ca 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,4 @@ -use anyhow::{anyhow, Context, Result}; +use anyhow::{anyhow, Result}; use chrono::Local; use clap::{Parser, Subcommand}; use colored::Colorize; @@ -7,35 +7,11 @@ use crossterm::{ execute, terminal::{Clear, ClearType}, }; +use lambda_cli::api::{LambdaClient, LambdaError}; use prettytable::{row, Table}; -use reqwest::blocking::Client; -use reqwest::header::AUTHORIZATION; -use serde::Deserialize; -use std::collections::HashMap; -use std::env; use std::io::{stdout, Write}; -use std::thread; use std::time::{Duration, Instant}; -use thiserror::Error; - -const API_BASE_URL: &str = "https://cloud.lambdalabs.com/api/v1"; -const DEFAULT_TIMEOUT_SECS: u64 = 30; - -#[derive(Error, Debug)] -pub enum LambdaError { - #[error("API key not set. Please set LAMBDA_API_KEY environment variable")] - ApiKeyNotSet, - #[error("Instance type '{0}' not found")] - InstanceTypeNotFound(String), - #[error("No regions available for instance type '{0}'")] - NoRegionsAvailable(String), - #[error("No instance IDs returned from launch request")] - NoInstanceIds, - #[error("API request failed: {0}")] - ApiError(String), - #[error("SSH key is required for this operation")] - SshKeyRequired, -} +use tokio::runtime::Runtime; /// A command-line tool for Lambda Labs cloud GPU API #[derive(Parser)] @@ -91,86 +67,6 @@ enum Commands { }, } -#[derive(Deserialize, Debug)] -struct ApiResponse { - data: T, -} - -#[derive(Deserialize, Debug)] -struct ApiErrorResponse { - error: ApiErrorDetail, -} - -#[derive(Deserialize, Debug)] -struct ApiErrorDetail { - message: String, -} - -#[derive(Deserialize, Debug)] -struct Instance { - id: Option, - name: Option, - status: Option, - ip: Option, - ssh_key_names: Option>, - instance_type: Option, - region: Option, -} - -#[derive(Deserialize, Debug)] -struct InstanceTypeInfo { - name: Option, -} - -#[derive(Deserialize, Debug)] -struct RegionInfo { - name: Option, -} - -#[derive(Deserialize, Debug)] -struct LaunchResponse { - instance_ids: Vec, -} - -#[derive(Deserialize, Debug, Clone)] -struct InstanceTypeResponse { - instance_type: InstanceType, - regions_with_capacity_available: Vec, -} - -#[derive(Deserialize, Debug, Clone)] -struct InstanceType { - description: String, - price_cents_per_hour: i32, - specs: InstanceSpecs, -} - -#[derive(Deserialize, Debug, Clone)] -struct InstanceSpecs { - vcpus: u32, - memory_gib: u32, - storage_gib: u32, -} - -#[derive(Deserialize, Debug, Clone)] -struct Region { - name: String, - #[allow(dead_code)] - description: String, -} - -fn create_client() -> Result { - Client::builder() - .timeout(Duration::from_secs(DEFAULT_TIMEOUT_SECS)) - .connect_timeout(Duration::from_secs(10)) - .build() - .context("Failed to create HTTP client") -} - -fn get_api_key() -> Result { - env::var("LAMBDA_API_KEY").map_err(|_| LambdaError::ApiKeyNotSet.into()) -} - fn main() { dotenv::dotenv().ok(); @@ -181,85 +77,38 @@ fn main() { } fn run() -> Result<()> { - // Parse CLI first so --help works without API key let cli = Cli::parse(); - - // Now get API key and client (needed for all commands) - let api_key = get_api_key()?; - let client = create_client()?; + let rt = Runtime::new()?; + let client = LambdaClient::from_env()?; match &cli.command { - Some(Commands::List) => list_instances(&client, &api_key), + Some(Commands::List) => list_instances(&rt, &client), Some(Commands::Start { gpu, ssh, name, region, - }) => start_instance( - &client, - &api_key, - gpu, - ssh, - name.as_deref(), - region.as_deref(), - ), - Some(Commands::Stop { instance_id }) => stop_instance(&client, &api_key, instance_id), - Some(Commands::Running) => list_running_instances(&client, &api_key), + }) => start_instance(&rt, &client, gpu, ssh, name.as_deref(), region.as_deref()), + Some(Commands::Stop { instance_id }) => stop_instance(&rt, &client, instance_id), + Some(Commands::Running) => list_running_instances(&rt, &client), Some(Commands::Find { gpu, ssh, interval, name, - }) => find_and_start_instance(&client, &api_key, gpu, ssh, *interval, name.as_deref()), - None => validate_api_key(&client, &api_key), + }) => find_and_start_instance(&rt, &client, gpu, ssh, *interval, name.as_deref()), + None => validate_api_key(&rt, &client), } } -fn validate_api_key(client: &Client, api_key: &str) -> Result<()> { - let url = format!("{}/instances", API_BASE_URL); - let response = client - .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", api_key)) - .send() - .context("Failed to connect to Lambda Labs API")?; - - if response.status().is_success() { - println!("{}", "API key is valid".green()); - Ok(()) - } else { - let status = response.status(); - let error_msg = parse_error_response(response); - Err(anyhow!( - "API key validation failed ({}): {}", - status, - error_msg - )) - } -} - -fn parse_error_response(response: reqwest::blocking::Response) -> String { - response - .json::() - .map(|e| e.error.message) - .unwrap_or_else(|_| "Unknown error".to_string()) +fn validate_api_key(rt: &Runtime, client: &LambdaClient) -> Result<()> { + rt.block_on(client.validate_api_key())?; + println!("{}", "API key is valid".green()); + Ok(()) } -fn list_instances(client: &Client, api_key: &str) -> Result<()> { - let url = format!("{}/instance-types", API_BASE_URL); - let response = client - .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", api_key)) - .send() - .context("Failed to fetch instance types")?; - - if !response.status().is_success() { - let error_msg = parse_error_response(response); - return Err(anyhow!("Failed to list instances: {}", error_msg)); - } - - let response: ApiResponse> = response - .json() - .context("Failed to parse instance types response")?; +fn list_instances(rt: &Runtime, client: &LambdaClient) -> Result<()> { + let types = rt.block_on(client.list_instance_types())?; let mut table = Table::new(); table.add_row(row![ @@ -272,38 +121,26 @@ fn list_instances(client: &Client, api_key: &str) -> Result<()> { "Available Regions" ]); - let mut entries: Vec<_> = response.data.into_iter().collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0)); - - for (key, instance_type_response) in entries { - let regions: Vec = instance_type_response - .regions_with_capacity_available - .iter() - .map(|region| region.name.clone()) - .collect(); - - let availability = if regions.is_empty() { + for t in types { + let availability = if t.regions_available.is_empty() { "None".red().to_string() } else { - regions.join(", ").blue().to_string() + t.regions_available.join(", ").blue().to_string() }; - let price = format!( - "${:.2}", - instance_type_response.instance_type.price_cents_per_hour as f64 / 100.0 - ); + let price = format!("${:.2}", t.price_cents_per_hour as f64 / 100.0); table.add_row(row![ - if regions.is_empty() { - key.dimmed().to_string() + if t.regions_available.is_empty() { + t.name.dimmed().to_string() } else { - key.green().to_string() + t.name.green().to_string() }, - instance_type_response.instance_type.description, + t.description, price.yellow(), - instance_type_response.instance_type.specs.vcpus, - instance_type_response.instance_type.specs.memory_gib, - instance_type_response.instance_type.specs.storage_gib, + t.vcpus, + t.memory_gib, + t.storage_gib, availability ]); } @@ -313,101 +150,33 @@ fn list_instances(client: &Client, api_key: &str) -> Result<()> { } fn start_instance( - client: &Client, - api_key: &str, + rt: &Runtime, + client: &LambdaClient, gpu: &str, ssh: &str, name: Option<&str>, region: Option<&str>, ) -> Result<()> { - let instance_type_response = get_instance_type_response(client, api_key, gpu)? - .ok_or_else(|| LambdaError::InstanceTypeNotFound(gpu.to_string()))?; - - let region_name = if let Some(r) = region { - // Validate the specified region is available - if !instance_type_response - .regions_with_capacity_available - .iter() - .any(|reg| reg.name == r) - { - return Err(anyhow!( - "Region '{}' is not available for instance type '{}'. Available regions: {}", - r, - gpu, - instance_type_response - .regions_with_capacity_available - .iter() - .map(|reg| reg.name.as_str()) - .collect::>() - .join(", ") - )); - } - r.to_string() - } else { - // Auto-select first available region - instance_type_response - .regions_with_capacity_available - .first() - .ok_or_else(|| LambdaError::NoRegionsAvailable(gpu.to_string()))? - .name - .clone() - }; - - let url = format!("{}/instance-operations/launch", API_BASE_URL); - - // Build payload - include name if provided - let mut payload = serde_json::json!({ - "region_name": region_name, - "instance_type_name": gpu, - "ssh_key_names": [ssh], - "quantity": 1 - }); - - if let Some(instance_name) = name { - payload["name"] = serde_json::Value::String(instance_name.to_string()); - } - println!( - "Launching {} in {}{}...", + "Launching {} {}...", gpu.green(), - region_name.blue(), - name.map(|n| format!(" as '{}'", n.cyan())) + name.map(|n| format!("as '{}'", n.cyan())) .unwrap_or_default() ); - let response = client - .post(&url) - .header(AUTHORIZATION, format!("Bearer {}", api_key)) - .json(&payload) - .send() - .context("Failed to send launch request")?; - - if !response.status().is_success() { - let error_msg = parse_error_response(response); - return Err(anyhow!("Failed to launch instance: {}", error_msg)); - } - - let response_text = response.text().context("Failed to read response")?; - let parsed_response: ApiResponse = - serde_json::from_str(&response_text).context("Failed to parse launch response")?; - - let instance_id = parsed_response - .data - .instance_ids - .first() - .ok_or(LambdaError::NoInstanceIds)?; + let result = rt.block_on(client.launch_instance(gpu, ssh, name, region))?; println!( "{} Instance {} launched in region {}", "Success!".green().bold(), - instance_id.cyan(), - region_name.blue() + result.instance_id.cyan(), + result.region.blue() ); println!("Waiting for instance to become active..."); - // Poll for instance to become ready instead of fixed sleep + // Poll for instance to become ready let start_time = Instant::now(); - let max_wait = Duration::from_secs(300); // 5 minutes max + let max_wait = Duration::from_secs(300); let poll_interval = Duration::from_secs(10); loop { @@ -419,9 +188,9 @@ fn start_instance( break; } - thread::sleep(poll_interval); + std::thread::sleep(poll_interval); - match get_instance_details(client, api_key, instance_id) { + match rt.block_on(client.get_instance(&result.instance_id)) { Ok(instance) => { let status = instance.status.as_deref().unwrap_or("unknown"); print!( @@ -452,7 +221,6 @@ fn start_instance( } } Err(e) => { - // Instance might not be queryable yet, continue polling print!("\r{} Waiting for instance... ", "Polling...".dimmed()); stdout().flush().ok(); eprintln!("\nWarning: {}", e); @@ -463,48 +231,10 @@ fn start_instance( Ok(()) } -fn get_instance_type_response( - client: &Client, - api_key: &str, - gpu: &str, -) -> Result> { - let url = format!("{}/instance-types", API_BASE_URL); - let response = client - .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", api_key)) - .send() - .context("Failed to fetch instance types")?; - - if !response.status().is_success() { - let error_msg = parse_error_response(response); - return Err(anyhow!("Failed to get instance types: {}", error_msg)); - } - - let response: ApiResponse> = - response.json().context("Failed to parse instance types")?; - - Ok(response.data.get(gpu).cloned()) -} - -fn stop_instance(client: &Client, api_key: &str, instance_id: &str) -> Result<()> { - let url = format!("{}/instance-operations/terminate", API_BASE_URL); - let payload = serde_json::json!({ - "instance_ids": [instance_id] - }); - +fn stop_instance(rt: &Runtime, client: &LambdaClient, instance_id: &str) -> Result<()> { println!("Terminating instance {}...", instance_id.cyan()); - let response = client - .post(&url) - .header(AUTHORIZATION, format!("Bearer {}", api_key)) - .json(&payload) - .send() - .context("Failed to send terminate request")?; - - if !response.status().is_success() { - let error_msg = parse_error_response(response); - return Err(anyhow!("Failed to terminate instance: {}", error_msg)); - } + rt.block_on(client.terminate_instance(instance_id))?; println!( "{} Instance {} terminated", @@ -514,24 +244,10 @@ fn stop_instance(client: &Client, api_key: &str, instance_id: &str) -> Result<() Ok(()) } -fn list_running_instances(client: &Client, api_key: &str) -> Result<()> { - let url = format!("{}/instances", API_BASE_URL); - let response = client - .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", api_key)) - .send() - .context("Failed to fetch running instances")?; - - if !response.status().is_success() { - let error_msg = parse_error_response(response); - return Err(anyhow!("Failed to list running instances: {}", error_msg)); - } - - let response: ApiResponse> = response - .json() - .context("Failed to parse running instances response")?; +fn list_running_instances(rt: &Runtime, client: &LambdaClient) -> Result<()> { + let instances = rt.block_on(client.list_running_instances())?; - if response.data.is_empty() { + if instances.is_empty() { println!("{}", "No running instances".yellow()); return Ok(()); } @@ -547,7 +263,7 @@ fn list_running_instances(client: &Client, api_key: &str) -> Result<()> { "SSH Keys" ]); - for instance in response.data { + for instance in instances { let status = instance.status.as_deref().unwrap_or("N/A"); let status_colored = match status { "active" => status.green().to_string(), @@ -581,8 +297,8 @@ fn list_running_instances(client: &Client, api_key: &str) -> Result<()> { } fn find_and_start_instance( - client: &Client, - api_key: &str, + rt: &Runtime, + client: &LambdaClient, gpu: &str, ssh: &str, interval: u64, @@ -599,45 +315,35 @@ fn find_and_start_instance( ); println!("Press Ctrl+C to stop\n"); - // Do first check immediately let mut first_check = true; loop { if !first_check { - thread::sleep(Duration::from_secs(interval)); + std::thread::sleep(Duration::from_secs(interval)); } first_check = false; let check_time = Local::now().format("%Y-%m-%d %H:%M:%S").to_string(); - match get_instance_type_response(client, api_key, gpu) { - Ok(Some(instance_type_response)) => { - if !instance_type_response - .regions_with_capacity_available - .is_empty() - { - let regions: Vec = instance_type_response - .regions_with_capacity_available - .iter() - .map(|region| region.name.clone()) - .collect(); - - // Clear screen and show success - execute!(stdout(), Clear(ClearType::All), MoveTo(0, 0)).ok(); - println!( - "{} Found {} available in: {}", - "SUCCESS!".green().bold(), - gpu.green(), - regions.join(", ").blue() - ); - - return start_instance(client, api_key, gpu, ssh, name, None); - } + match rt.block_on(client.check_availability(gpu)) { + Ok(regions) if !regions.is_empty() => { + execute!(stdout(), Clear(ClearType::All), MoveTo(0, 0)).ok(); + println!( + "{} Found {} available in: {}", + "SUCCESS!".green().bold(), + gpu.green(), + regions.join(", ").blue() + ); + + return start_instance(rt, client, gpu, ssh, name, None); } - Ok(None) => { - return Err(LambdaError::InstanceTypeNotFound(gpu.to_string()).into()); + Ok(_) => { + // No availability } Err(e) => { + if e.to_string().contains("not found") { + return Err(LambdaError::InstanceTypeNotFound(gpu.to_string()).into()); + } eprintln!( "{} Failed to check availability: {}", "Warning:".yellow(), @@ -647,7 +353,6 @@ fn find_and_start_instance( } } - // Update status display execute!(stdout(), Clear(ClearType::All), MoveTo(0, 0)).ok(); let mut table = Table::new(); table.add_row(row!["Instance Type", "Last Checked", "Status"]); @@ -656,299 +361,3 @@ fn find_and_start_instance( println!("\nNext check in {} seconds... (Ctrl+C to stop)", interval); } } - -fn get_instance_details(client: &Client, api_key: &str, instance_id: &str) -> Result { - let url = format!("{}/instances/{}", API_BASE_URL, instance_id); - let response = client - .get(&url) - .header(AUTHORIZATION, format!("Bearer {}", api_key)) - .send() - .context("Failed to fetch instance details")?; - - if !response.status().is_success() { - let status = response.status(); - let error_msg = parse_error_response(response); - return Err(anyhow!( - "Failed to get instance details ({}): {}", - status, - error_msg - )); - } - - let response: ApiResponse = response - .json() - .context("Failed to parse instance details")?; - - Ok(response.data) -} - -#[cfg(test)] -mod tests { - use super::*; - - // Test API response parsing - #[test] - fn test_parse_instance_response() { - let json = r#"{ - "data": { - "id": "inst-123", - "name": "my-instance", - "status": "active", - "ip": "192.168.1.1", - "ssh_key_names": ["key1", "key2"], - "instance_type": {"name": "gpu_1x_a100"}, - "region": {"name": "us-west-1"} - } - }"#; - - let response: ApiResponse = serde_json::from_str(json).unwrap(); - assert_eq!(response.data.id, Some("inst-123".to_string())); - assert_eq!(response.data.name, Some("my-instance".to_string())); - assert_eq!(response.data.status, Some("active".to_string())); - assert_eq!(response.data.ip, Some("192.168.1.1".to_string())); - assert_eq!( - response.data.ssh_key_names, - Some(vec!["key1".to_string(), "key2".to_string()]) - ); - } - - #[test] - fn test_parse_instance_response_with_nulls() { - let json = r#"{ - "data": { - "id": "inst-123", - "status": "booting" - } - }"#; - - let response: ApiResponse = serde_json::from_str(json).unwrap(); - assert_eq!(response.data.id, Some("inst-123".to_string())); - assert_eq!(response.data.name, None); - assert_eq!(response.data.ip, None); - assert_eq!(response.data.ssh_key_names, None); - } - - #[test] - fn test_parse_launch_response() { - let json = r#"{ - "data": { - "instance_ids": ["inst-abc", "inst-def"] - } - }"#; - - let response: ApiResponse = serde_json::from_str(json).unwrap(); - assert_eq!(response.data.instance_ids.len(), 2); - assert_eq!(response.data.instance_ids[0], "inst-abc"); - assert_eq!(response.data.instance_ids[1], "inst-def"); - } - - #[test] - fn test_parse_instance_types_response() { - let json = r#"{ - "data": { - "gpu_1x_a100": { - "instance_type": { - "description": "1x A100 GPU", - "price_cents_per_hour": 110, - "specs": { - "vcpus": 24, - "memory_gib": 200, - "storage_gib": 512 - } - }, - "regions_with_capacity_available": [ - {"name": "us-west-1", "description": "US West"}, - {"name": "us-east-1", "description": "US East"} - ] - } - } - }"#; - - let response: ApiResponse> = - serde_json::from_str(json).unwrap(); - - assert!(response.data.contains_key("gpu_1x_a100")); - let instance_type = response.data.get("gpu_1x_a100").unwrap(); - assert_eq!(instance_type.instance_type.description, "1x A100 GPU"); - assert_eq!(instance_type.instance_type.price_cents_per_hour, 110); - assert_eq!(instance_type.instance_type.specs.vcpus, 24); - assert_eq!(instance_type.instance_type.specs.memory_gib, 200); - assert_eq!(instance_type.regions_with_capacity_available.len(), 2); - assert_eq!( - instance_type.regions_with_capacity_available[0].name, - "us-west-1" - ); - } - - #[test] - fn test_parse_empty_regions() { - let json = r#"{ - "data": { - "gpu_8x_h100": { - "instance_type": { - "description": "8x H100 GPU", - "price_cents_per_hour": 2400, - "specs": { - "vcpus": 208, - "memory_gib": 1800, - "storage_gib": 20000 - } - }, - "regions_with_capacity_available": [] - } - } - }"#; - - let response: ApiResponse> = - serde_json::from_str(json).unwrap(); - - let instance_type = response.data.get("gpu_8x_h100").unwrap(); - assert!(instance_type.regions_with_capacity_available.is_empty()); - } - - #[test] - fn test_parse_api_error_response() { - let json = r#"{ - "error": { - "message": "Invalid API key" - } - }"#; - - let response: ApiErrorResponse = serde_json::from_str(json).unwrap(); - assert_eq!(response.error.message, "Invalid API key"); - } - - // Test error types - #[test] - fn test_lambda_error_messages() { - assert_eq!( - LambdaError::ApiKeyNotSet.to_string(), - "API key not set. Please set LAMBDA_API_KEY environment variable" - ); - assert_eq!( - LambdaError::InstanceTypeNotFound("gpu_1x_a100".to_string()).to_string(), - "Instance type 'gpu_1x_a100' not found" - ); - assert_eq!( - LambdaError::NoRegionsAvailable("gpu_1x_a100".to_string()).to_string(), - "No regions available for instance type 'gpu_1x_a100'" - ); - assert_eq!( - LambdaError::NoInstanceIds.to_string(), - "No instance IDs returned from launch request" - ); - assert_eq!( - LambdaError::SshKeyRequired.to_string(), - "SSH key is required for this operation" - ); - } - - // Test URL construction - #[test] - fn test_api_base_url() { - assert_eq!(API_BASE_URL, "https://cloud.lambdalabs.com/api/v1"); - } - - #[test] - fn test_url_construction() { - let instance_id = "inst-123"; - let url = format!("{}/instances/{}", API_BASE_URL, instance_id); - assert_eq!( - url, - "https://cloud.lambdalabs.com/api/v1/instances/inst-123" - ); - } - - // Test client creation - #[test] - fn test_create_client() { - let client = create_client(); - assert!(client.is_ok()); - } - - // Test JSON payload construction - #[test] - fn test_launch_payload_without_name() { - let payload = serde_json::json!({ - "region_name": "us-west-1", - "instance_type_name": "gpu_1x_a100", - "ssh_key_names": ["my-key"], - "quantity": 1 - }); - - assert_eq!(payload["region_name"], "us-west-1"); - assert_eq!(payload["instance_type_name"], "gpu_1x_a100"); - assert_eq!(payload["quantity"], 1); - assert!(payload.get("name").is_none()); - } - - #[test] - fn test_launch_payload_with_name() { - let mut payload = serde_json::json!({ - "region_name": "us-west-1", - "instance_type_name": "gpu_1x_a100", - "ssh_key_names": ["my-key"], - "quantity": 1 - }); - - let instance_name = "my-training-job"; - payload["name"] = serde_json::Value::String(instance_name.to_string()); - - assert_eq!(payload["name"], "my-training-job"); - } - - #[test] - fn test_terminate_payload() { - let instance_id = "inst-123"; - let payload = serde_json::json!({ - "instance_ids": [instance_id] - }); - - let ids = payload["instance_ids"].as_array().unwrap(); - assert_eq!(ids.len(), 1); - assert_eq!(ids[0], "inst-123"); - } - - // Test InstanceTypeResponse cloning (needed for the lookup function) - #[test] - fn test_instance_type_response_clone() { - let response = InstanceTypeResponse { - instance_type: InstanceType { - description: "Test GPU".to_string(), - price_cents_per_hour: 100, - specs: InstanceSpecs { - vcpus: 8, - memory_gib: 64, - storage_gib: 256, - }, - }, - regions_with_capacity_available: vec![Region { - name: "us-west-1".to_string(), - description: "US West".to_string(), - }], - }; - - let cloned = response.clone(); - assert_eq!(cloned.instance_type.description, "Test GPU"); - assert_eq!(cloned.regions_with_capacity_available.len(), 1); - } - - // Test safe array access patterns - #[test] - fn test_first_on_empty_vec() { - let empty: Vec = vec![]; - assert!(empty.first().is_none()); - } - - #[test] - fn test_first_on_non_empty_vec() { - let vec = vec!["first".to_string(), "second".to_string()]; - assert_eq!(vec.first(), Some(&"first".to_string())); - } - - // Test default timeout value - #[test] - fn test_default_timeout() { - assert_eq!(DEFAULT_TIMEOUT_SECS, 30); - } -}