From 35d6ace552332173dd8171b16fd5552dce6842b8 Mon Sep 17 00:00:00 2001 From: xodapi <4956501+xodapi@users.noreply.github.com> Date: Sat, 4 Jul 2026 21:50:57 +0700 Subject: [PATCH] feat(mcp): add stdio server for IDE integration (closes #12) --- Cargo.lock | 134 ++++++++++++- Cargo.toml | 2 + src/cli/args.rs | 4 + src/cli/cache.rs | 55 +++++- src/cli/mcp.rs | 498 +++++++++++++++++++++++++++++++++++++++++++++++ src/cli/mod.rs | 1 + src/main.rs | 5 + 7 files changed, 697 insertions(+), 2 deletions(-) create mode 100644 src/cli/mcp.rs diff --git a/Cargo.lock b/Cargo.lock index d31b458..77d1bfa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -395,7 +395,7 @@ dependencies = [ "log", "num-rational", "num-traits", - "pastey", + "pastey 0.1.1", "rayon", "thiserror 2.0.18", "v_frame", @@ -764,6 +764,7 @@ dependencies = [ "iana-time-zone", "js-sys", "num-traits", + "serde", "wasm-bindgen", "windows-link", ] @@ -1377,6 +1378,12 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "edf234dd1594d6dd434a8fb8cada51ddbbc593e40e4a01556a0b31c62da2775b" +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "ed25519" version = "2.2.3" @@ -4180,6 +4187,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35fb2e5f958ec131621fdd531e9fc186ed768cbe395337403ae56c17a74c68ec" +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -4781,6 +4794,26 @@ dependencies = [ "thiserror 2.0.18", ] +[[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 2.0.118", +] + [[package]] name = "regex" version = "1.12.4" @@ -4890,6 +4923,41 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rmcp" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d1f571c72940a19d9532fe52dbea8bc9912bf1d766c2970bb824056b86f3f59" +dependencies = [ + "async-trait", + "base64", + "chrono", + "futures", + "pastey 0.2.3", + "pin-project-lite", + "rmcp-macros", + "schemars", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "rmcp-macros" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aad0035b69380782d78ea95b508327e6deaa2235909053e596eea8f27b5e1d5" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "serde_json", + "syn 2.0.118", +] + [[package]] name = "rowan" version = "0.16.1" @@ -5054,6 +5122,32 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "schemars" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" +dependencies = [ + "chrono", + "dyn-clone", + "ref-cast", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d115b50f4aaeea07e79c1912f645c7513d81715d0420f8bc77a18c6260b307f" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn 2.0.118", +] + [[package]] name = "scoped-tls" version = "1.0.1" @@ -5171,6 +5265,17 @@ dependencies = [ "syn 2.0.118", ] +[[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", + "syn 2.0.118", +] + [[package]] name = "serde_json" version = "1.0.150" @@ -6015,9 +6120,21 @@ dependencies = [ "mio", "pin-project-lite", "socket2", + "tokio-macros", "windows-sys 0.61.2", ] +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -6028,6 +6145,19 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + [[package]] name = "toml" version = "0.5.11" @@ -6495,12 +6625,14 @@ dependencies = [ "ratatui", "redb", "reqwest", + "rmcp", "self_update", "serde", "serde_json", "sha2 0.11.0", "slint", "slint-build", + "tokio", "toml 0.8.2", "tray-icon", "winres", diff --git a/Cargo.toml b/Cargo.toml index f2b8b79..ff7def1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -54,6 +54,7 @@ crossterm = "0.28" ratatui = { version = "0.28", default-features = false, features = ["crossterm"] } reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } redb = "2" +rmcp = { version = "1.7.0", features = ["server", "schemars", "transport-io"] } serde = { version = "1", features = ["derive"] } serde_json = "1" jni = { version = "0.22.4", optional = true } @@ -64,6 +65,7 @@ rand = "0.9" self_update = { version = "0.41", default-features = false, features = ["archive-tar", "archive-zip", "compression-flate2", "compression-zip-deflate", "rustls"] } sha2 = "0.11.0" fs3 = "0.5.0" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "io-std"] } [dev-dependencies] insta = "1.48" diff --git a/src/cli/args.rs b/src/cli/args.rs index 8c3fd8b..f361372 100644 --- a/src/cli/args.rs +++ b/src/cli/args.rs @@ -49,6 +49,7 @@ pub struct Args { pub list_accounts: bool, pub doctor: bool, pub init: bool, + pub mcp: bool, pub vpn: bool, pub auto_failover: bool, pub no_cache: bool, @@ -88,6 +89,7 @@ where list_accounts: false, doctor: false, init: false, + mcp: false, vpn: false, auto_failover: true, no_cache: false, @@ -122,6 +124,7 @@ where "--list-accounts" => parsed.list_accounts = true, "--doctor" => parsed.doctor = true, "--init" => parsed.init = true, + "--mcp" => parsed.mcp = true, "--vpn" => parsed.vpn = true, "--no-failover" => parsed.auto_failover = false, "--no-cache" => parsed.no_cache = true, @@ -320,6 +323,7 @@ OPTIONS: --list-accounts List available account profiles --doctor Run system diagnostics --init Interactive setup wizard + --mcp Start MCP server on stdio --trend Show 30-day usage trends (requires saved snapshots) --days Days of trend history to show [default: 30] diff --git a/src/cli/cache.rs b/src/cli/cache.rs index 6a02b07..762bc73 100644 --- a/src/cli/cache.rs +++ b/src/cli/cache.rs @@ -1,4 +1,4 @@ -use redb::{Database, TableDefinition}; +use redb::{Database, ReadableTable, TableDefinition}; use serde_json::Value; use std::path::PathBuf; use std::time::{Duration, Instant}; @@ -6,6 +6,14 @@ use std::time::{Duration, Instant}; const TABLE: TableDefinition<&[u8], &[u8]> = TableDefinition::new("api_cache"); const DEFAULT_TTL_SECS: u64 = 30; +pub struct CacheEntry { + pub cache_key: String, + pub api_base: String, + pub cached_at: String, + pub age_secs: u64, + pub payload: Value, +} + pub struct CacheStore { db: Database, ttl: Duration, @@ -112,6 +120,51 @@ impl CacheStore { pub fn ttl(&self) -> Duration { self.ttl } + + pub fn entries(&self) -> Result, String> { + let tx = self + .db + .begin_read() + .map_err(|e| format!("cache read tx failed: {e}"))?; + let table = tx + .open_table(TABLE) + .map_err(|e| format!("cache table open failed: {e}"))?; + let now_secs = std::time::UNIX_EPOCH + .elapsed() + .map(|d| d.as_secs()) + .unwrap_or(0); + let mut entries = Vec::new(); + for item in table + .iter() + .map_err(|e| format!("cache iteration failed: {e}"))? + { + let (key, value) = item.map_err(|e| format!("cache row read failed: {e}"))?; + let cache_key = std::str::from_utf8(key.value()) + .map_err(|e| format!("cache key is invalid UTF-8: {e}"))? + .to_string(); + let parsed: Value = serde_json::from_slice(value.value()) + .map_err(|e| format!("cache value is invalid JSON: {e}"))?; + let cached_at = parsed + .get("cached_at") + .and_then(Value::as_str) + .unwrap_or("0") + .to_string(); + let cached_at_secs = cached_at.parse::().unwrap_or(0); + let payload = parsed.get("payload").cloned().unwrap_or(Value::Null); + let api_base = cache_key + .split_once('|') + .map(|(_, api_base)| api_base.to_string()) + .unwrap_or_default(); + entries.push(CacheEntry { + cache_key, + api_base, + cached_at, + age_secs: now_secs.saturating_sub(cached_at_secs), + payload, + }); + } + Ok(entries) + } } fn cache_key(api_key: &str, api_base: &str) -> String { diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs new file mode 100644 index 0000000..175af0d --- /dev/null +++ b/src/cli/mcp.rs @@ -0,0 +1,498 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use rmcp::{ + ErrorData as McpError, Json, ServerHandler, ServiceExt, + handler::server::{router::tool::ToolRouter, wrapper::Parameters}, + model::{ + AnnotateAble, GetPromptRequestParams, GetPromptResult, ListPromptsResult, + ListResourceTemplatesResult, ListResourcesResult, PaginatedRequestParams, Prompt, + PromptArgument, PromptMessage, PromptMessageRole, RawResource, RawResourceTemplate, + ReadResourceRequestParams, ReadResourceResult, ResourceContents, ServerCapabilities, + ServerInfo, + }, + schemars, tool, tool_handler, tool_router, +}; +use serde_json::{Value, json}; + +use crate::{self as ng, cli}; + +#[derive(Debug, Clone, serde::Deserialize, schemars::JsonSchema)] +struct CheckLimitsRequest { + #[serde(default)] + account: Option, + #[serde(default)] + no_cache: bool, +} + +#[derive(Debug, Clone, serde::Deserialize, schemars::JsonSchema)] +struct InvalidateCacheRequest { + #[serde(default)] + account: Option, +} + +#[derive(Debug, Clone, serde::Serialize, schemars::JsonSchema)] +struct ToolJsonObject { + #[serde(flatten)] + data: HashMap, +} + +impl ToolJsonObject { + fn from_value(value: Value) -> Result { + let data = value + .as_object() + .cloned() + .ok_or_else(|| "expected JSON object".to_string())? + .into_iter() + .collect(); + Ok(Self { data }) + } +} + +#[derive(Debug, Clone)] +pub struct VimitMcpServer { + shared: Arc>, + tool_router: ToolRouter, +} + +#[derive(Debug)] +struct SharedState { + account_names: Vec, + current_account_idx: usize, +} + +impl VimitMcpServer { + pub fn new(account_names: Vec) -> Self { + Self { + shared: Arc::new(Mutex::new(SharedState { + account_names, + current_account_idx: 0, + })), + tool_router: Self::tool_router(), + } + } + + fn selected_account_name(&self) -> Option { + let state = self.shared.lock().unwrap(); + state.account_names.get(state.current_account_idx).cloned() + } + + fn switch_account_name(&self) -> Option { + let mut state = self.shared.lock().unwrap(); + if state.account_names.is_empty() { + return None; + } + state.current_account_idx = (state.current_account_idx + 1) % state.account_names.len(); + state.account_names.get(state.current_account_idx).cloned() + } + + fn runtime_args(&self, no_cache: bool) -> cli::args::Args { + cli::args::Args { + api_base: None, + api_key_env: cli::constants::DEFAULT_API_KEY_ENV.to_string(), + env_file: None, + demo: false, + mock: None, + output: cli::args::OutputMode::Json, + monitor: false, + preset: cli::args::Preset::Full, + theme: cli::theme::Theme::Btop, + with_abtop: false, + notify: false, + watch: 0, + fail_on: cli::args::FailOn::Never, + warning_threshold: cli::constants::DEFAULT_WARNING_THRESHOLD, + danger_threshold: cli::constants::DEFAULT_DANGER_THRESHOLD, + window_thresholds: HashMap::new(), + account: self.selected_account_name(), + list_accounts: false, + doctor: false, + init: false, + mcp: true, + vpn: false, + auto_failover: true, + no_cache, + trend: false, + trend_days: 30, + update: false, + update_check: false, + help: false, + version: false, + config: None, + daily_limit: None, + overlay: false, + } + } + + fn load_runtime( + &self, + requested_account: Option<&str>, + no_cache: bool, + ) -> Result { + let args = self.runtime_args(no_cache); + let config = cli::config::Config::load(None)?; + let accounts = cli::accounts::AccountsConfig::load()?; + + let selected_name = requested_account + .map(ToOwned::to_owned) + .or_else(|| self.selected_account_name()); + let selected_account = match selected_name.as_deref() { + Some(name) => Some(accounts.resolve(name)?), + None => None, + }; + + let api_base = selected_account + .as_ref() + .and_then(|account| account.api_base.clone()) + .or_else(|| config.api_base.clone()); + let api_key_env = selected_account + .as_ref() + .and_then(|account| account.api_key_env.clone()) + .unwrap_or_else(|| args.api_key_env.clone()); + + let mut runtime = ng::RuntimeConfig::from_dotenv(api_base, &api_key_env, None)?; + runtime.auto_failover = config.auto_failover.unwrap_or(true); + + Ok(RuntimeContext { + args, + runtime, + account_name: selected_name, + }) + } + + fn status_json( + &self, + requested_account: Option<&str>, + no_cache: bool, + ) -> Result { + let ctx = self.load_runtime(requested_account, no_cache)?; + let http = ng::HttpClient::new(ng::USER_AGENT)?; + let cache = cli::cache::CacheStore::open()?; + let mut router = ng::Router::new( + ctx.runtime.api_base.clone(), + ng::api_fallbacks_for(&ctx.runtime.api_base, ctx.runtime.auto_failover), + ); + let mut daily_file = crate::cli::daily::DailyFile::load(); + let snapshot = cli::monitor::collect_status( + &ctx.args, + &ctx.runtime, + &http, + cache.as_ref(), + Some(&mut router), + &mut daily_file, + )?; + let mut status = ng::summary_to_json_with_stale( + &snapshot.windows, + snapshot.abtop.as_ref(), + snapshot.daily.as_ref(), + snapshot.stale, + snapshot.latency_ms, + &snapshot.api_endpoint, + ); + if let Some(map) = status.as_object_mut() { + map.insert( + "account".to_string(), + ctx.account_name + .map(Value::String) + .unwrap_or(Value::String("default".to_string())), + ); + } + Ok(status) + } + + fn cache_json(&self, requested_account: Option<&str>) -> Result { + let ctx = self.load_runtime(requested_account, false)?; + let cache = cli::cache::CacheStore::open()?; + let entries = match cache { + Some(store) => store.entries()?, + None => Vec::new(), + }; + let entries: Vec = entries + .into_iter() + .filter(|entry| entry.api_base == ctx.runtime.api_base) + .map(|entry| { + json!({ + "cache_key": entry.cache_key, + "api_base": entry.api_base, + "cached_at": entry.cached_at, + "age_secs": entry.age_secs, + "payload": entry.payload, + }) + }) + .collect(); + Ok(json!({ + "account": ctx.account_name.unwrap_or_else(|| "default".to_string()), + "api_base": ctx.runtime.api_base, + "entries": entries, + })) + } + + fn trends_json(&self, days: u64) -> Result { + let store = cli::trends::TrendStore::open()?; + let trends = match store { + Some(store) => store.query_trends(days)?, + None => Vec::new(), + }; + let trends: Vec = trends + .into_iter() + .map(|day| { + json!({ + "date": day.date.to_string(), + "windows": day.windows.into_iter().map(|window| { + json!({ + "key": window.key, + "samples": window.samples, + "peak_max": window.peak_max, + "peak_avg": window.peak_avg, + "credits_avg_used": window.credits_avg_used, + "credits_avg_limit": window.credits_avg_limit, + "requests_avg_used": window.requests_avg_used, + "requests_avg_limit": window.requests_avg_limit, + }) + }).collect::>() + }) + }) + .collect(); + Ok(json!({ + "days": days, + "trends": trends, + })) + } +} + +impl Default for VimitMcpServer { + fn default() -> Self { + let accounts = cli::accounts::AccountsConfig::load() + .map(|accounts| accounts.list_names()) + .unwrap_or_default(); + Self::new(accounts) + } +} + +struct RuntimeContext { + args: cli::args::Args, + runtime: ng::RuntimeConfig, + account_name: Option, +} + +#[tool_router(router = tool_router)] +impl VimitMcpServer { + #[tool(description = "Run a status check and return JSON matching the CLI JSON output")] + async fn check_limits( + &self, + Parameters(request): Parameters, + ) -> Result, String> { + self.status_json(request.account.as_deref(), request.no_cache) + .and_then(ToolJsonObject::from_value) + .map(Json) + } + + #[tool(description = "Invalidate the cache entry for the selected or requested account")] + async fn invalidate_cache( + &self, + Parameters(request): Parameters, + ) -> Result, String> { + let ctx = self.load_runtime(request.account.as_deref(), false)?; + let cache = cli::cache::CacheStore::open()?; + if let Some(store) = cache { + store.remove(&ctx.runtime.api_key, &ctx.runtime.api_base)?; + } + ToolJsonObject::from_value(json!({ + "ok": true, + "account": ctx.account_name.unwrap_or_else(|| "default".to_string()), + "api_base": ctx.runtime.api_base, + })) + .map(Json) + } + + #[tool( + description = "Rotate to the next configured account and return the selected account name" + )] + async fn switch_account(&self) -> Result, String> { + ToolJsonObject::from_value(json!({ + "account": self.switch_account_name().unwrap_or_else(|| "default".to_string()), + })) + .map(Json) + } +} + +#[tool_handler(router = self.tool_router)] +impl ServerHandler for VimitMcpServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new( + ServerCapabilities::builder() + .enable_tools() + .enable_resources() + .enable_prompts() + .build(), + ) + .with_instructions( + "vimit MCP exposes current quota status, cache contents, trend history, and account helpers.", + ) + } + + fn list_resources( + &self, + _request: Option, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + '_ { + std::future::ready({ + Ok(ListResourcesResult { + resources: vec![ + RawResource::new("vimit://status/current", "Current quota status") + .with_description("Current quota status snapshot in vimit JSON form") + .no_annotation(), + RawResource::new("vimit://cache/current", "Current cache entries") + .with_description("Cached API payloads for the active account") + .no_annotation(), + RawResource::new("vimit://trends/30d", "30 day trends") + .with_description("Historical trend samples for the last 30 days") + .no_annotation(), + ], + ..Default::default() + }) + }) + } + + fn list_resource_templates( + &self, + _request: Option, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + '_ { + std::future::ready({ + Ok(ListResourceTemplatesResult { + resource_templates: vec![ + RawResourceTemplate::new("vimit://trends/{days}", "Trend history") + .with_description( + "Historical trend samples for the requested number of days", + ) + .no_annotation(), + ], + ..Default::default() + }) + }) + } + + fn read_resource( + &self, + request: ReadResourceRequestParams, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + '_ { + std::future::ready((|| { + let value = if request.uri == "vimit://status/current" { + self.status_json(None, false) + } else if request.uri == "vimit://cache/current" { + self.cache_json(None) + } else if let Some(days) = request.uri.strip_prefix("vimit://trends/") { + days.parse::() + .map_err(|error| format!("invalid trend day value '{days}': {error}")) + .and_then(|days| self.trends_json(days)) + } else { + Err(format!("unknown resource: {}", request.uri)) + } + .map_err(|error| McpError::resource_not_found(error, None))?; + + let text = serde_json::to_string_pretty(&value) + .map_err(|error| McpError::internal_error(error.to_string(), None))?; + Ok(ReadResourceResult::new(vec![ResourceContents::text( + text, + request.uri, + )])) + })()) + } + + fn list_prompts( + &self, + _request: Option, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + '_ { + std::future::ready({ + Ok(ListPromptsResult { + prompts: vec![Prompt::new( + "usage_summary", + Some("Summarize current vimit quota state for review or planning context."), + Some(vec![ + PromptArgument::new("focus") + .with_description("Optional focus such as review, planning, or release") + .with_required(false), + ]), + )], + ..Default::default() + }) + }) + } + + fn get_prompt( + &self, + request: GetPromptRequestParams, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + '_ { + std::future::ready((|| { + if request.name != "usage_summary" { + return Err(McpError::invalid_params( + format!("unknown prompt: {}", request.name), + None, + )); + } + + let focus = request + .arguments + .as_ref() + .and_then(|args| args.get("focus")) + .and_then(Value::as_str) + .unwrap_or("general"); + let snapshot = self + .status_json(None, false) + .map_err(|error| McpError::internal_error(error, None))?; + let snapshot = serde_json::to_string_pretty(&snapshot) + .map_err(|error| McpError::internal_error(error.to_string(), None))?; + + Ok(GetPromptResult::new(vec![PromptMessage::new_text( + PromptMessageRole::User, + format!( + "Use this vimit quota snapshot to produce a concise {focus} summary.\n\n{snapshot}" + ), + )])) + })()) + } +} + +pub fn run_mcp() -> Result { + let runtime = tokio::runtime::Runtime::new() + .map_err(|error| format!("cannot initialize tokio runtime for MCP: {error}"))?; + runtime.block_on(async move { + let accounts = cli::accounts::AccountsConfig::load()?; + let server = VimitMcpServer::new(accounts.list_names()); + server + .serve((tokio::io::stdin(), tokio::io::stdout())) + .await + .map_err(|error| error.to_string())? + .waiting() + .await + .map_err(|error| error.to_string())?; + Ok(0) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn switch_account_cycles_through_configured_accounts() { + let server = VimitMcpServer::new(vec!["alpha".into(), "beta".into(), "gamma".into()]); + assert_eq!(server.selected_account_name().as_deref(), Some("alpha")); + assert_eq!(server.switch_account_name().as_deref(), Some("beta")); + assert_eq!(server.switch_account_name().as_deref(), Some("gamma")); + assert_eq!(server.switch_account_name().as_deref(), Some("alpha")); + } + + #[test] + fn runtime_args_enable_mcp_mode() { + let server = VimitMcpServer::new(Vec::new()); + let args = server.runtime_args(true); + assert!(args.mcp); + assert!(args.no_cache); + assert!(matches!(args.output, cli::args::OutputMode::Json)); + } +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 2e73b32..b53314a 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -6,6 +6,7 @@ pub mod constants; pub mod daily; pub mod doctor; pub mod init; +pub mod mcp; pub mod monitor; pub mod notify; pub mod output; diff --git a/src/main.rs b/src/main.rs index df211d9..02836ad 100644 --- a/src/main.rs +++ b/src/main.rs @@ -118,6 +118,10 @@ fn real_main() -> Result { return cli::init::run_init(); } + if cli_args.mcp { + return cli::mcp::run_mcp(); + } + let mut merged = config.merge_with_defaults()?; if let Some(ref account_name) = cli_args.account { let account = accounts.resolve(account_name)?; @@ -311,6 +315,7 @@ fn merge_args_with_config(args: Args, merged: &MergedConfig) -> Args { list_accounts: args.list_accounts, doctor: args.doctor, init: args.init, + mcp: args.mcp, vpn: args.vpn, auto_failover: args.auto_failover && merged.auto_failover, no_cache: args.no_cache,