| 1 | |
|
| 2 | |
|
| 3 | |
|
| 4 | |
|
| 5 | |
|
| 6 | |
|
| 7 | |
|
| 8 | |
|
| 9 | | use anyhow::{Context, Result}; |
| 10 | | use serde_json::{json, Value as Json}; |
| 11 | | use spectra_core::keys::KeyPair; |
| 12 | | use spectra_mcp::client::Client; |
| 13 | | use std::path::{Path, PathBuf}; |
| 14 | | use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; |
| 15 | |
|
| 16 | |
|
| 17 | | #[derive(Clone, Debug)] |
| 18 | | pub struct Label { |
| 19 | | pub name: String, |
| 20 | | pub kind: String, |
| 21 | | pub folder: Option<PathBuf>, |
| 22 | | pub home: PathBuf, |
| 23 | | } |
| 24 | |
|
| 25 | 0× | pub fn agents_dir() -> PathBuf { |
| 26 | 0× | if let Some(d) = std::env::var_os("SPECTRA_AGENTS").filter(|d| !d.is_empty()) { |
| 27 | 0× | return PathBuf::from(d); |
| 28 | 0× | } |
| 29 | 0× | std::env::var_os("HOME").map(PathBuf::from).unwrap_or_default().join(".spectra-agents") |
| 30 | 0× | } |
| 31 | |
|
| 32 | 26× | pub fn write_label(home: &Path, name: &str, kind: &str, project: &str, folder: Option<&str>) -> Result<()> { |
| 33 | 26× | let mut t = toml::Table::new(); |
| 34 | 26× | t.insert("name".into(), name.into()); |
| 35 | 26× | t.insert("kind".into(), kind.into()); |
| 36 | 26× | t.insert("project".into(), project.into()); |
| 37 | 26× | if let Some(f) = folder.filter(|f| !f.is_empty()) { |
| 38 | 24× | t.insert("folder".into(), f.into()); |
| 39 | 24× | } |
| 40 | 26× | let text = format!("# Written by `spectra agent add`: lets `spectra mcp --auto` find this agent.\n{t}"); |
| 41 | 26× | spectra_core::durable::write_atomic(&home.join("agent.toml"), text.as_bytes())?; |
| 42 | 26× | Ok(()) |
| 43 | 26× | } |
| 44 | |
|
| 45 | 0× | fn read_label(home: &Path) -> Option<Label> { |
| 46 | 0× | let t: toml::Table = std::fs::read_to_string(home.join("agent.toml")).ok()?.parse().ok()?; |
| 47 | 0× | let s = |k: &str| t.get(k).and_then(|v| v.as_str()).map(String::from); |
| 48 | 0× | s("project")?; |
| 49 | 0× | Some(Label { name: s("name")?, kind: s("kind")?, folder: s("folder").map(PathBuf::from), home: home.to_path_buf() }) |
| 50 | 0× | } |
| 51 | |
|
| 52 | |
|
| 53 | |
|
| 54 | 0× | async fn learn_label(endpoint: &spectra_mcp::client::Endpoint, home: &Path) -> Option<Label> { |
| 55 | 0× | let key = KeyPair::load_home(home).ok()?; |
| 56 | 0× | let c = Client::at(endpoint.clone(), Some(key.clone())); |
| 57 | 0× | let agents = c.send("GET", "/api/v1/agents", &[], vec![]).await.ok()?.json().ok()?; |
| 58 | 0× | let me = agents.as_array()?.iter().find(|a| a["key"] == key.key.to_hex())?.clone(); |
| 59 | 0× | let projects = c.send("GET", "/api/v1/projects", &[], vec![]).await.ok()?.json().ok()?; |
| 60 | 0× | let project = me["project"].as_str()?.to_string(); |
| 61 | 0× | let folder = projects.as_array()?.iter().find(|p| p["name"] == project.as_str()).and_then(|p| p["path"].as_str().map(String::from)); |
| 62 | 0× | let (name, kind) = (me["name"].as_str()?.to_string(), me["kind"].as_str()?.to_string()); |
| 63 | 0× | let _ = write_label(home, &name, &kind, &project, folder.as_deref()); |
| 64 | 0× | Some(Label { name, kind, folder: folder.map(PathBuf::from), home: home.to_path_buf() }) |
| 65 | 0× | } |
| 66 | |
|
| 67 | |
|
| 68 | 0× | pub async fn labels(endpoint: Option<&spectra_mcp::client::Endpoint>) -> Vec<Label> { |
| 69 | 0× | let mut out = Vec::new(); |
| 70 | 0× | let Ok(dirs) = std::fs::read_dir(agents_dir()) else { return out }; |
| 71 | 0× | for d in dirs.filter_map(|e| e.ok()).map(|e| e.path()).filter(|p| p.join("keys").is_dir()) { |
| 72 | 0× | match read_label(&d) { |
| 73 | 0× | Some(l) => out.push(l), |
| 74 | | None => { |
| 75 | 0× | if let Some(ep) = endpoint { |
| 76 | 0× | if let Some(l) = learn_label(ep, &d).await { |
| 77 | 0× | out.push(l); |
| 78 | 0× | } |
| 79 | 0× | } |
| 80 | | } |
| 81 | | } |
| 82 | | } |
| 83 | 0× | out |
| 84 | 0× | } |
| 85 | |
|
| 86 | |
|
| 87 | |
|
| 88 | 5× | pub fn choose<'a>(labels: &'a [Label], cwd: &Path, kind: &str) -> Option<&'a Label> { |
| 89 | 5× | let cwd = cwd.canonicalize().unwrap_or_else(|_| cwd.to_path_buf()); |
| 90 | 5× | labels |
| 91 | 5× | .iter() |
| 92 | 15× | .filter(|l| l.kind == kind) |
| 93 | 9× | .filter_map(|l| { |
| 94 | 9× | let f = l.folder.as_ref()?; |
| 95 | 9× | let f = f.canonicalize().unwrap_or_else(|_| f.clone()); |
| 96 | 9× | cwd.starts_with(&f).then(|| (f.components().count(), l)) |
| 97 | 9× | }) |
| 98 | 5× | .max_by(|(a, x), (b, y)| a.cmp(b).then_with(|| y.name.cmp(&x.name))) |
| 99 | 5× | .map(|(_, l)| l) |
| 100 | 5× | } |
| 101 | |
|
| 102 | |
|
| 103 | |
|
| 104 | 0× | pub async fn unconnected(cwd: &Path, kind: &str) -> Result<i32> { |
| 105 | 0× | let why = format!( |
| 106 | | "This folder ({}) has no Spectra agent of kind {kind}, so Spectra's tools are not available here. \ |
| 107 | | The owner can add one with the Spectra VS Code extension (Spectra: Add Agent) or with \ |
| 108 | | `spectra agent add NAME --kind {kind} --project PROJECT --home ~/.spectra-agents/NAME`, \ |
| 109 | | after `spectra project add {}`; then start a new session.", |
| 110 | 0× | cwd.display(), |
| 111 | 0× | cwd.display() |
| 112 | | ); |
| 113 | 0× | let mut out = tokio::io::stdout(); |
| 114 | 0× | let mut lines = BufReader::new(tokio::io::stdin()).lines(); |
| 115 | 0× | while let Some(line) = lines.next_line().await? { |
| 116 | 0× | let Ok(req) = serde_json::from_str::<Json>(&line) else { continue }; |
| 117 | 0× | let Some(id) = req.get("id").cloned() else { continue }; |
| 118 | 0× | let result = match req.get("method").and_then(Json::as_str).unwrap_or("") { |
| 119 | 0× | "initialize" => { |
| 120 | 0× | let v = req.pointer("/params/protocolVersion").and_then(Json::as_str).unwrap_or(super::mcp::STDIO_VERSIONS[0]); |
| 121 | 0× | let v = if super::mcp::STDIO_VERSIONS.contains(&v) { v } else { super::mcp::STDIO_VERSIONS[0] }; |
| 122 | 0× | json!({ "protocolVersion": v, "capabilities": { "tools": {} }, "serverInfo": { "name": "spectra", "version": env!("CARGO_PKG_VERSION") }, "instructions": why }) |
| 123 | | } |
| 124 | 0× | "tools/list" => json!({ "tools": [] }), |
| 125 | 0× | "ping" => json!({}), |
| 126 | | _ => { |
| 127 | 0× | let e = json!({ "jsonrpc": "2.0", "id": id, "error": { "code": -32601, "message": why } }); |
| 128 | 0× | out.write_all(format!("{e}\n").as_bytes()).await?; |
| 129 | 0× | out.flush().await?; |
| 130 | 0× | continue; |
| 131 | | } |
| 132 | | }; |
| 133 | 0× | out.write_all(format!("{}\n", json!({ "jsonrpc": "2.0", "id": id, "result": result })).as_bytes()).await?; |
| 134 | 0× | out.flush().await?; |
| 135 | | } |
| 136 | 0× | Ok(0) |
| 137 | 0× | } |
| 138 | |
|
| 139 | |
|
| 140 | 0× | pub async fn resolve(ctx: &super::Ctx, kind: &str) -> Result<Option<Label>> { |
| 141 | 0× | let cwd = std::env::current_dir().context("this session's folder")?; |
| 142 | 0× | let endpoint = ctx.endpoint().ok(); |
| 143 | 0× | let all = labels(endpoint.as_ref()).await; |
| 144 | 0× | Ok(choose(&all, &cwd, kind).cloned()) |
| 145 | 0× | } |
| 146 | |
|
| 147 | | #[cfg(test)] |
| 148 | | mod tests { |
| 149 | | use super::*; |
| 150 | |
|
| 151 | 3× | fn l(name: &str, kind: &str, folder: &str) -> Label { |
| 152 | 3× | Label { name: name.into(), kind: kind.into(), folder: Some(PathBuf::from(folder)), home: PathBuf::new() } |
| 153 | 3× | } |
| 154 | |
|
| 155 | | #[test] |
| 156 | 1× | fn the_deepest_folder_of_the_right_kind_wins() { |
| 157 | 1× | let all = [ |
| 158 | 1× | l("claude-dev", "claude-code", "/dev"), |
| 159 | 1× | l("claude-book", "claude-code", "/dev/book"), |
| 160 | 1× | l("kimi-book", "kimi-code", "/dev/book"), |
| 161 | 1× | ]; |
| 162 | 1× | assert_eq!(choose(&all, Path::new("/dev/book/ch3"), "claude-code").unwrap().name, "claude-book"); |
| 163 | 1× | assert_eq!(choose(&all, Path::new("/dev/other"), "claude-code").unwrap().name, "claude-dev"); |
| 164 | 1× | assert_eq!(choose(&all, Path::new("/dev/book"), "kimi-code").unwrap().name, "kimi-book"); |
| 165 | 1× | assert!(choose(&all, Path::new("/elsewhere"), "claude-code").is_none()); |
| 166 | 1× | assert!(choose(&all, Path::new("/devx"), "claude-code").is_none(), "a prefix of the name is not inside"); |
| 167 | 1× | } |
| 168 | | } |