| 1 | //! SplitFire's [`Profile`] for `ed-acp`: who the agent is, its prompt and tools, its slash |
| 2 | //! commands, and the optional built-in `demucs --mcp` stem separator. |
| 3 | |
| 4 | use std::path::{Path, PathBuf}; |
| 5 | use std::sync::Arc; |
| 6 | use std::time::Duration; |
| 7 | |
| 8 | use agent_client_protocol::schema::v1::{AvailableCommand, PromptCapabilities}; |
| 9 | use async_trait::async_trait; |
| 10 | use ed_acp::ed_mcp::{McpTool, McpToolset, ServerSpec, SharedServer, find_binary, switched_off}; |
| 11 | use ed_acp::{ |
| 12 | AgentInfo, AudioHints, LlmEnv, McpCall, Profile, PromptCtx, SessionCtx, Toolset, |
| 13 | WorkspaceTools, llm, |
| 14 | }; |
| 15 | use serde_json::{Value, json}; |
| 16 | |
| 17 | use crate::prompt; |
| 18 | use crate::tools::MusicTools; |
| 19 | |
| 20 | pub const INFO: AgentInfo = AgentInfo { |
| 21 | name: "splitfire-agent", |
| 22 | display_name: "SplitFire", |
| 23 | version: env!("CARGO_PKG_VERSION"), |
| 24 | env: LlmEnv { |
| 25 | prefix: "SPLITFIRE", |
| 26 | dir_name: "splitfire-agent", |
| 27 | default_model: llm::DEFAULT_MODEL, |
| 28 | }, |
| 29 | }; |
| 30 | |
| 31 | /// Server name of the built-in stem separator; a client-supplied server of this name wins. |
| 32 | pub const DEMUCS: &str = "demucs"; |
| 33 | |
| 34 | /// `demucs --mcp`, started once per process on first need: `SPLITFIRE_DEMUCS_BIN`, else |
| 35 | /// `demucs` on `PATH`. `SPLITFIRE_MCP_DEMUCS=off` turns it off. |
| 36 | static DEMUCS_SERVER: SharedServer = SharedServer::new(DEMUCS, demucs_spec); |
| 37 | |
| 38 | fn demucs_spec() -> Option<ServerSpec> { |
| 39 | if switched_off("SPLITFIRE_MCP_DEMUCS") { |
| 40 | tracing::info!("stem separation disabled (SPLITFIRE_MCP_DEMUCS=off)"); |
| 41 | return None; |
| 42 | } |
| 43 | let Some(command) = find_binary("SPLITFIRE_DEMUCS_BIN", "demucs") else { |
| 44 | tracing::info!( |
| 45 | "stem separation unavailable: `demucs` not found (install it or set SPLITFIRE_DEMUCS_BIN)" |
| 46 | ); |
| 47 | return None; |
| 48 | }; |
| 49 | Some(ServerSpec::stdio( |
| 50 | DEMUCS, |
| 51 | command, |
| 52 | vec!["--mcp".to_string()], |
| 53 | Vec::new(), |
| 54 | )) |
| 55 | } |
| 56 | |
| 57 | /// Whether a stem separator (built-in or client-supplied) is connected. |
| 58 | fn has_stems(mcp: &McpToolset) -> bool { |
| 59 | mcp.has_tool(DEMUCS, "separate_stems") |
| 60 | } |
| 61 | |
| 62 | fn is_separation(tool: &McpTool) -> bool { |
| 63 | tool.server == DEMUCS && tool.tool == "separate_stems" |
| 64 | } |
| 65 | |
| 66 | pub struct SplitFire; |
| 67 | |
| 68 | #[async_trait] |
| 69 | impl Profile for SplitFire { |
| 70 | fn info(&self) -> AgentInfo { |
| 71 | INFO |
| 72 | } |
| 73 | |
| 74 | fn system_prompt(&self, ctx: &PromptCtx<'_>) -> String { |
| 75 | prompt::system_prompt(ctx.surface, ctx.cwd, ctx.roots, has_stems(ctx.mcp)) |
| 76 | } |
| 77 | |
| 78 | fn toolsets(&self) -> Vec<Arc<dyn Toolset>> { |
| 79 | vec![Arc::new(WorkspaceTools), Arc::new(MusicTools)] |
| 80 | } |
| 81 | |
| 82 | fn slash_commands(&self, ctx: &SessionCtx<'_>) -> Vec<AvailableCommand> { |
| 83 | prompt::slash_commands(has_stems(ctx.mcp)) |
| 84 | } |
| 85 | |
| 86 | fn expand_slash(&self, text: &str, ctx: &SessionCtx<'_>) -> Option<String> { |
| 87 | prompt::expand_slash(text, has_stems(ctx.mcp)) |
| 88 | } |
| 89 | |
| 90 | fn builtin_mcp_servers(&self) -> Vec<&'static SharedServer> { |
| 91 | vec![&DEMUCS_SERVER] |
| 92 | } |
| 93 | |
| 94 | fn prompt_capabilities(&self) -> PromptCapabilities { |
| 95 | PromptCapabilities::new() |
| 96 | .embedded_context(true) |
| 97 | .image(true) |
| 98 | .audio(true) |
| 99 | } |
| 100 | |
| 101 | fn audio_hints(&self) -> AudioHints { |
| 102 | AudioHints { |
| 103 | link: |path| { |
| 104 | format!( |
| 105 | "[Referenced audio file: {}. Use analyze_audio on that path to inspect it.]", |
| 106 | path.display() |
| 107 | ) |
| 108 | }, |
| 109 | attachment: |mime, path| { |
| 110 | format!( |
| 111 | "[Attached audio ({mime}) saved at {}. Use analyze_audio on that path to inspect it.]", |
| 112 | path.display() |
| 113 | ) |
| 114 | }, |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | /// For stem separation, say what will be separated and whether the model weights still |
| 119 | /// have to be downloaded, per the server's `list_models`. |
| 120 | async fn mcp_permission_preview(&self, call: McpCall<'_>) -> Option<String> { |
| 121 | if !is_separation(call.tool) { |
| 122 | return None; |
| 123 | } |
| 124 | let input = call |
| 125 | .args |
| 126 | .get("input") |
| 127 | .and_then(Value::as_str) |
| 128 | .unwrap_or("?"); |
| 129 | let model = call |
| 130 | .args |
| 131 | .get("model") |
| 132 | .and_then(Value::as_str) |
| 133 | .unwrap_or("htdemucs"); |
| 134 | let mut text = |
| 135 | format!("Separate {input} into stems with {model}. Runs locally and can take minutes."); |
| 136 | let note = match call.connection { |
| 137 | Some(conn) => { |
| 138 | let result = conn |
| 139 | .call_raw("list_models", json!({}), Duration::from_secs(10)) |
| 140 | .await; |
| 141 | result.and_then(|r| download_note(&r, model)) |
| 142 | } |
| 143 | None => None, |
| 144 | }; |
| 145 | text.push('\n'); |
| 146 | text.push_str( |
| 147 | note.as_deref() |
| 148 | .unwrap_or("Could not check whether the model weights are cached."), |
| 149 | ); |
| 150 | Some(text) |
| 151 | } |
| 152 | |
| 153 | fn mcp_result_locations(&self, tool: &McpTool, structured: &Value) -> Vec<PathBuf> { |
| 154 | if is_separation(tool) { |
| 155 | stem_paths(structured) |
| 156 | } else { |
| 157 | Vec::new() |
| 158 | } |
| 159 | } |
| 160 | } |
| 161 | |
| 162 | /// "Downloads N MB…" when `model` is not cached, from a `list_models` result. |
| 163 | fn download_note(result: &Value, model: &str) -> Option<String> { |
| 164 | let models = result.pointer("/structuredContent/models")?.as_array()?; |
| 165 | let m = models |
| 166 | .iter() |
| 167 | .find(|m| m.get("id").and_then(Value::as_str) == Some(model))?; |
| 168 | let size = m.get("size_mb").and_then(Value::as_u64); |
| 169 | Some(match (m.get("cached").and_then(Value::as_bool), size) { |
| 170 | (Some(true), _) => "The model weights are already downloaded.".to_string(), |
| 171 | (_, Some(mb)) => format!("First use downloads about {mb} MB of model weights."), |
| 172 | _ => "The model weights may need to be downloaded first.".to_string(), |
| 173 | }) |
| 174 | } |
| 175 | |
| 176 | /// `structuredContent.stems[].path`, absolute paths only. |
| 177 | fn stem_paths(structured: &Value) -> Vec<PathBuf> { |
| 178 | structured |
| 179 | .get("stems") |
| 180 | .and_then(Value::as_array) |
| 181 | .map(|stems| { |
| 182 | stems |
| 183 | .iter() |
| 184 | .filter_map(|s| s.get("path").and_then(Value::as_str)) |
| 185 | .map(PathBuf::from) |
| 186 | .filter(|p| Path::new(p).is_absolute()) |
| 187 | .collect() |
| 188 | }) |
| 189 | .unwrap_or_default() |
| 190 | } |
| 191 | |
| 192 | #[cfg(test)] |
| 193 | mod tests { |
| 194 | use super::*; |
| 195 | |
| 196 | #[test] |
| 197 | fn stems_become_absolute_locations() { |
| 198 | let v = json!({"stems": [ |
| 199 | {"id": "drums", "path": "/tmp/a/drums.wav"}, |
| 200 | {"id": "bass", "path": "relative/bass.wav"}, |
| 201 | {"id": "x"} |
| 202 | ]}); |
| 203 | let paths = stem_paths(&v); |
| 204 | assert_eq!(paths, vec![PathBuf::from("/tmp/a/drums.wav")]); |
| 205 | assert!(stem_paths(&json!({})).is_empty()); |
| 206 | } |
| 207 | |
| 208 | #[test] |
| 209 | fn download_note_reads_list_models() { |
| 210 | let r = json!({"structuredContent": {"models": [ |
| 211 | {"id": "htdemucs", "size_mb": 80, "cached": false}, |
| 212 | {"id": "htdemucs_ft", "size_mb": 320, "cached": true} |
| 213 | ]}}); |
| 214 | assert_eq!( |
| 215 | download_note(&r, "htdemucs").as_deref(), |
| 216 | Some("First use downloads about 80 MB of model weights.") |
| 217 | ); |
| 218 | assert_eq!( |
| 219 | download_note(&r, "htdemucs_ft").as_deref(), |
| 220 | Some("The model weights are already downloaded.") |
| 221 | ); |
| 222 | assert_eq!(download_note(&r, "nope"), None); |
| 223 | } |
| 224 | } |