@splitfireai / agent / commits / 73aee25

Refactor: remove agent/llm/mcp/tui/session modules, add profile

- Delete src/agent.rs, src/llm.rs, src/mcp.rs, src/session_store.rs, src/tools/workspace.rs, src/tui.rs - Add src/profile.rs - Trim main.rs, prompt.rs, tools/mod.rs, mcp_stdio_stub.rs down to leaner implementations - Update Cargo.toml and Cargo.lock dependencies accordingly - Update tests (acp_compliance, common) to match new structure - Update README and registry/agent.json metadata Co-Authored-By: siGit Code v1.6.5-acp <noreply@sigit.si>

Seto committed Oct 9, 2026 at 12:12 UTC 73aee2588702fbe743b7ef787e8bb2e499bb7a04
18 files changed +478 -6236
Cargo.lock
+99 -7
@@ -174,6 +174,17 @@ version = "4.7.1"
174 source = "registry+https://github.com/rust-lang/crates.io-index"
175 checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de"
176
177 +[[package]]
178 +name = "async-trait"
179 +version = "0.1.92"
180 +source = "registry+https://github.com/rust-lang/crates.io-index"
181 +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
182 +dependencies = [
183 + "proc-macro2",
184 + "quote",
185 + "syn 3.0.6",
186 +]
187 +
188 [[package]]
189 name = "atomic-waker"
190 version = "1.1.2"
@@ -622,6 +633,71 @@ dependencies = [
633 "smallvec",
634 ]
635
636 +[[package]]
637 +name = "ed-acp"
638 +version = "1.1.0"
639 +source = "registry+https://github.com/rust-lang/crates.io-index"
640 +checksum = "3acecd5520efb32a7bebd6d02b865de3d785ba4ebd32ede8a4f595ae4e5e9e79"
641 +dependencies = [
642 + "agent-client-protocol",
643 + "anyhow",
644 + "async-trait",
645 + "base64",
646 + "directories",
647 + "ed-mcp",
648 + "futures",
649 + "reqwest",
650 + "serde",
651 + "serde_json",
652 + "tokio",
653 + "tokio-util",
654 + "tracing",
655 + "tracing-subscriber",
656 + "uuid",
657 +]
658 +
659 +[[package]]
660 +name = "ed-acp-testkit"
661 +version = "1.1.0"
662 +source = "registry+https://github.com/rust-lang/crates.io-index"
663 +checksum = "e9a99f0cd1c4c72a2235edc42f03c5b8a7b0d18e29d0064c467a1120423af252"
664 +dependencies = [
665 + "serde_json",
666 + "tokio",
667 + "uuid",
668 +]
669 +
670 +[[package]]
671 +name = "ed-acp-tui"
672 +version = "1.1.0"
673 +source = "registry+https://github.com/rust-lang/crates.io-index"
674 +checksum = "1241d2bc5655c8c880f3dbd57746d29cbc25fe616869346583d73ffbf8821b73"
675 +dependencies = [
676 + "agent-client-protocol",
677 + "anyhow",
678 + "crossterm",
679 + "futures",
680 + "ratatui",
681 + "tokio",
682 +]
683 +
684 +[[package]]
685 +name = "ed-mcp"
686 +version = "1.1.0"
687 +source = "registry+https://github.com/rust-lang/crates.io-index"
688 +checksum = "578adbef751312c45b43fdd8615da51bc45adc7ddf039f9dd2030369e2b07ff6"
689 +dependencies = [
690 + "agent-client-protocol",
691 + "anyhow",
692 + "futures",
693 + "http",
694 + "rmcp",
695 + "serde_json",
696 + "tokio",
697 + "tokio-util",
698 + "tracing",
699 +]
700 +
701 [[package]]
702 name = "either"
703 version = "1.18.0"
@@ -2010,13 +2086,18 @@ version = "3.5.1"
2086 source = "registry+https://github.com/rust-lang/crates.io-index"
2087 checksum = "ab5023804cbe0304cbe06197ac16e645a0a0d0ffa2953bf735c15c2fa0213350"
2088 dependencies = [
2089 + "base64",
2090 + "bytes",
2091 "chrono",
2092 "futures",
2093 + "http",
2094 "indexmap 2.14.2",
2095 "pin-project-lite",
2096 "process-wrap",
2097 + "reqwest",
2098 "serde",
2099 "serde_json",
2100 + "sse-stream",
2101 "thiserror",
2102 "tokio",
2103 "tokio-stream",
@@ -2427,26 +2508,37 @@ version = "1.0.0"
2508 dependencies = [
2509 "agent-client-protocol",
2510 "anyhow",
2511 + "async-trait",
2512 "base64",
2431 - "crossterm",
2513 "directories",
2514 "ebur128",
2434 - "futures",
2515 + "ed-acp",
2516 + "ed-acp-testkit",
2517 + "ed-acp-tui",
2518 + "ed-mcp",
2519 "hound",
2436 - "ratatui",
2520 "realfft",
2438 - "reqwest",
2439 - "rmcp",
2521 "serde",
2522 "serde_json",
2523 "symphonia",
2524 "tokio",
2444 - "tokio-util",
2525 "tracing",
2446 - "tracing-subscriber",
2526 "uuid",
2527 ]
2528
2529 +[[package]]
2530 +name = "sse-stream"
2531 +version = "0.2.6"
2532 +source = "registry+https://github.com/rust-lang/crates.io-index"
2533 +checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4"
2534 +dependencies = [
2535 + "bytes",
2536 + "futures-util",
2537 + "http-body",
2538 + "http-body-util",
2539 + "pin-project-lite",
2540 +]
2541 +
2542 [[package]]
2543 name = "stable_deref_trait"
2544 version = "1.2.1"
Cargo.toml
+7 -8
@@ -7,25 +7,24 @@ description = "Music-domain ACP agent: theory, production, industry, culture and
7 repository = "https://github.com/splitfireai/splitfire-agent"
8
9 [dependencies]
10 +# Onde Agent Platform: the ACP server, terminal UI and MCP layer, from the onde-ed workspace.
11 +ed-acp = { version = "1.1.0" }
12 +ed-acp-tui = { version = "1.1.0" }
13 +ed-mcp = { version = "1.1.0", features = ["test-server"] }
14 agent-client-protocol = "2.2"
15 anyhow = "1.0.104"
12 -base64 = "0.23.1"
13 -crossterm = { version = "0.29.0", features = ["event-stream"] }
16 +async-trait = "0.1"
17 directories = "6.0.0"
18 ebur128 = "0.1.10"
16 -futures = "0.3.34"
17 -ratatui = { version = "0.30.2", default-features = false, features = ["crossterm", "layout-cache", "unstable-rendered-line-info"] }
19 realfft = "3.5.0"
19 -reqwest = { version = "0.13", default-features = false, features = ["rustls", "json", "stream"] }
20 -rmcp = { version = "3.5.1", default-features = false, features = ["client", "transport-child-process"] }
20 serde = { version = "1.0.229", features = ["derive"] }
21 serde_json = "1.0.151"
22 symphonia = { version = "0.6.1", features = ["mp3", "flac", "ogg", "vorbis", "wav", "pcm", "aac", "isomp4", "aiff", "alac", "mkv"] }
23 tokio = { version = "1.53.1", features = ["full"] }
25 -tokio-util = "0.7.19"
24 tracing = "0.1.44"
27 -tracing-subscriber = { version = "0.3.23", features = ["env-filter"] }
25 uuid = { version = "1.27.0", features = ["v4"] }
26
27 [dev-dependencies]
28 +ed-acp-testkit = { version = "1.1.0" }
29 +base64 = "0.23.1"
30 hound = "3.5.1"
LICENSE
+1 -1
@@ -186,7 +186,7 @@
186 same "printed page" as the copyright notice for easier
187 identification within third-party archives.
188
189 - Copyright 2026 OndeCode contributors
189 + Copyright 2026 SplitFire AI
190
191 Licensed under the Apache License, Version 2.0 (the "License");
192 you may not use this file except in compliance with the License.
README.md
+13 -11
@@ -139,21 +139,23 @@ cargo clippy --all-targets -- -D warnings
139 cargo test --locked
140 ```
141
142 -The integration tests run the real binary against a scripted mock model, plus a stub MCP server
143 -(`src/bin/mcp_stdio_stub.rs`). `tests/live.rs` talks to Onde Cloud and is skipped unless
144 -`ONDE_API_KEY` is set.
142 +SplitFire is built on the Onde Agent Platform crates in the
143 +[`onde-ed`](https://github.com/ondeinference/ed) workspace, checked out next to this repository:
144 +`ed-acp` (the ACP server: sessions, auth, the turn loop, approval, workspace tools),
145 +`ed-acp-tui` (the terminal UI) and `ed-mcp` (MCP over `rmcp`). This crate holds the music.
146 +
147 +The integration tests run the real binary against a scripted mock model and a test MCP server
148 +(`src/bin/mcp_stdio_stub.rs`, from `ed-mcp`). `tests/acp_compliance.rs` runs the platform's ACP
149 +conformance suite from `ed-acp-testkit` plus SplitFire's own tests. `tests/live.rs` talks to Onde
150 +Cloud and is skipped unless `ONDE_API_KEY` is set.
151
152 ```
147 -src/main.rs flags, mode dispatch, ACP handlers
148 -src/agent.rs sessions, turn loop, persistence
149 -src/prompt.rs system prompt, slash commands, prompt content
150 -src/llm.rs streaming Onde Cloud client
151 -src/session_store.rs on-disk sessions
152 -src/mcp.rs MCP client (client-supplied servers + demucs)
153 -src/tools/ workspace, theory and audio tools
153 +src/main.rs flags and mode dispatch (ACP, TUI, --setup, --list-models)
154 +src/profile.rs the ed-acp Profile: identity, tools, commands, built-in demucs
155 +src/prompt.rs system prompt and slash commands
156 +src/tools/ theory, audio analysis and separation-library tools
157 src/theory/ pitch, interval, scale, chord, key, tempo
158 src/audio/ decoding and analysis
156 -src/tui.rs terminal UI
159 registry/ ACP registry submission
160 ```
161
registry/agent.json
+2 -2
@@ -1,8 +1,8 @@
1 {
2 "id": "splitfire-agent",
3 - "name": "SplitFire",
3 + "name": "SplitFire AI",
4 "version": "1.0.0",
5 - "description": "Music assistant for musicians, producers and engineers: exact music theory tools (scales, chords, keys, tempo maths), audio analysis (loudness, tempo and key estimates), optional stem separation, and project file tools. Backed by Onde Inference; ships its own TUI and speaks ACP over stdio.",
5 + "description": "Music assistant for musicians, producers and engineers: exact music theory tools (scales, chords, keys, tempo maths), audio analysis (loudness, tempo and key estimates), optional stem separation, and project file tools.",
6 "repository": "https://github.com/splitfireai/splitfire-agent",
7 "authors": [
8 "SplitFire AI"
src/agent.rs deleted
-1237
@@ -1,1237 +0,0 @@
1 -//! `SplitFireAgent`: the ACP v1 agent. Owns sessions (in memory, backed by
2 -//! [`SessionStore`]), the turn loop, and the config-option model picker.
3 -
4 -use std::collections::{HashMap, HashSet};
5 -use std::path::{Path, PathBuf};
6 -use std::sync::atomic::{AtomicBool, Ordering};
7 -use std::sync::{Arc, Mutex};
8 -use std::time::{SystemTime, UNIX_EPOCH};
9 -
10 -use agent_client_protocol::schema::ProtocolVersion;
11 -use agent_client_protocol::schema::v1::{
12 - AuthMethod, AuthMethodTerminal, AuthenticateRequest, AuthenticateResponse,
13 - AvailableCommandsUpdate, ClientCapabilities, CloseSessionRequest, CloseSessionResponse,
14 - ConfigOptionUpdate, ContentBlock, ContentChunk, DeleteSessionRequest, DeleteSessionResponse,
15 - ListSessionsRequest, ListSessionsResponse, LoadSessionRequest, LoadSessionResponse,
16 - LogoutResponse, McpServer, MessageId, NewSessionRequest, NewSessionResponse, PromptRequest,
17 - PromptResponse, ResumeSessionRequest, ResumeSessionResponse, SessionConfigOption,
18 - SessionConfigOptionCategory, SessionConfigSelectOption, SessionConfigSelectOptions, SessionId,
19 - SessionInfo, SessionInfoUpdate, SessionNotification, SessionUpdate,
20 - SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, StopReason, ToolCall,
21 - ToolCallContent, ToolCallLocation, ToolCallStatus, ToolCallUpdateFields, UsageUpdate,
22 -};
23 -use agent_client_protocol::{Client, ConnectionTo, Responder};
24 -use serde_json::{Value, json};
25 -use tokio_util::sync::CancellationToken;
26 -
27 -use crate::llm::{Delta, LlmClient, LlmConfig};
28 -use crate::mcp::{self, McpToolset};
29 -use crate::prompt;
30 -use crate::session_store::{SessionStore, StoredSession};
31 -use crate::tools::{self, ToolCtx, ToolOutcome, absolutize, normalize};
32 -
33 -const MAX_TURNS: usize = 50;
34 -/// Id of the session config option that selects the model.
35 -pub const MODEL_CONFIG_ID: &str = "model";
36 -/// Id of the boolean session config option that skips permission prompts.
37 -pub const AUTO_APPROVE_CONFIG_ID: &str = "auto_approve";
38 -/// Context window reported in `usage_update` when SPLITFIRE_CONTEXT_WINDOW is unset.
39 -const DEFAULT_CONTEXT_WINDOW: u64 = 128_000;
40 -/// Id of the terminal auth method advertised on initialize.
41 -pub const AUTH_METHOD_ID: &str = "terminal-setup";
42 -/// Cap on the sessions returned by `session/list`.
43 -const MAX_LISTED: usize = 200;
44 -
45 -/// The terminal auth method: clients run `splitfire-agent --setup`.
46 -pub fn auth_methods() -> Vec<AuthMethod> {
47 - vec![AuthMethod::Terminal(
48 - AuthMethodTerminal::new(AUTH_METHOD_ID, "Run in terminal")
49 - .description("Interactive setup: store your Onde Inference API key")
50 - .args(vec!["--setup".into()]),
51 - )]
52 -}
53 -
54 -pub fn auth_required_error() -> agent_client_protocol::Error {
55 - agent_client_protocol::Error::auth_required().data(
56 - "No Onde API key configured. Authenticate with the terminal method (`splitfire-agent --setup`).",
57 - )
58 -}
59 -
60 -/// The only protocol version this agent speaks is v1. A client proposing a newer version gets
61 -/// v1 back (it can then decide whether to continue); an older (v0) client is told v1 too, as
62 -/// the spec requires the latest supported version when the requested one isn't supported.
63 -pub fn negotiate_version(client: ProtocolVersion) -> ProtocolVersion {
64 - let _ = client;
65 - ProtocolVersion::V1
66 -}
67 -
68 -struct Session {
69 - cwd: PathBuf,
70 - /// Additional workspace roots (from the request's `additional_directories`).
71 - roots: Vec<PathBuf>,
72 - model: String,
73 - /// Conversation history without the system prompt (rebuilt every turn).
74 - messages: Vec<Value>,
75 - cancel: CancellationToken,
76 - always_allowed: Arc<Mutex<HashSet<String>>>,
77 - always_rejected: Arc<Mutex<HashSet<String>>>,
78 - /// Approve tool calls without asking; the `auto_approve` config option. Not persisted:
79 - /// a reopened session starts from the process default (`--yolo`).
80 - auto_approve: Arc<AtomicBool>,
81 - title: Option<String>,
82 - /// Last activity, seconds since the Unix epoch.
83 - updated_at: u64,
84 - mcp: Arc<McpToolset>,
85 - /// A turn is running; its history is checked out of `messages`.
86 - busy: bool,
87 -}
88 -
89 -impl Session {
90 - fn fresh(cwd: PathBuf, roots: Vec<PathBuf>, model: String, yolo: bool) -> Self {
91 - Self::restored(
92 - StoredSession {
93 - cwd,
94 - roots,
95 - model,
96 - title: None,
97 - updated_at: now_secs(),
98 - messages: Vec::new(),
99 - },
100 - yolo,
101 - )
102 - }
103 -
104 - /// A session brought back from the store. Permission grants start fresh; MCP servers are
105 - /// connected by the request that reopens it.
106 - fn restored(stored: StoredSession, yolo: bool) -> Self {
107 - Self {
108 - cwd: stored.cwd,
109 - roots: stored.roots,
110 - model: stored.model,
111 - messages: stored.messages,
112 - cancel: CancellationToken::new(),
113 - always_allowed: Arc::default(),
114 - always_rejected: Arc::default(),
115 - auto_approve: Arc::new(AtomicBool::new(yolo)),
116 - title: stored.title,
117 - updated_at: stored.updated_at,
118 - mcp: Arc::new(McpToolset::default()),
119 - busy: false,
120 - }
121 - }
122 -
123 - fn stored(&self) -> StoredSession {
124 - StoredSession {
125 - cwd: self.cwd.clone(),
126 - roots: self.roots.clone(),
127 - model: self.model.clone(),
128 - title: self.title.clone(),
129 - updated_at: self.updated_at,
130 - messages: self.messages.clone(),
131 - }
132 - }
133 -}
134 -
135 -fn new_message_id() -> MessageId {
136 - MessageId::new(uuid::Uuid::new_v4().to_string())
137 -}
138 -
139 -/// ACP v1 requires absolute `cwd` and `additionalDirectories`.
140 -fn check_absolute(cwd: &Path, roots: &[PathBuf]) -> agent_client_protocol::Result<()> {
141 - if !cwd.is_absolute() {
142 - return Err(invalid("cwd must be an absolute path"));
143 - }
144 - if roots.iter().any(|r| !r.is_absolute()) {
145 - return Err(invalid(
146 - "additionalDirectories entries must be absolute paths",
147 - ));
148 - }
149 - Ok(())
150 -}
151 -
152 -pub fn now_secs() -> u64 {
153 - SystemTime::now()
154 - .duration_since(UNIX_EPOCH)
155 - .map(|d| d.as_secs())
156 - .unwrap_or(0)
157 -}
158 -
159 -/// ISO 8601 UTC timestamp without sub-second precision (no chrono dependency).
160 -pub fn iso8601(secs: u64) -> String {
161 - let days = secs / 86_400;
162 - let secs_of_day = secs % 86_400;
163 - let (h, m, s) = (
164 - secs_of_day / 3600,
165 - secs_of_day % 3600 / 60,
166 - secs_of_day % 60,
167 - );
168 - // Civil-from-days algorithm (Howard Hinnant).
169 - let z = days as i64 + 719_468;
170 - let era = z.div_euclid(146_097);
171 - let doe = z.rem_euclid(146_097);
172 - let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
173 - let y = yoe + era * 400;
174 - let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
175 - let mp = (5 * doy + 2) / 153;
176 - let d = doy - (153 * mp + 2) / 5 + 1;
177 - let mo = if mp < 10 { mp + 3 } else { mp - 9 };
178 - let y = if mo <= 2 { y + 1 } else { y };
179 - format!("{y:04}-{mo:02}-{d:02}T{h:02}:{m:02}:{s:02}Z")
180 -}
181 -
182 -/// Derive a short session title from the first user prompt.
183 -pub fn session_title(prompt: &str) -> String {
184 - const MAX: usize = 80;
185 - let first = prompt.lines().next().unwrap_or("").trim();
186 - if first.chars().count() <= MAX {
187 - first.to_string()
188 - } else {
189 - let cut: String = first.chars().take(MAX - 1).collect();
190 - format!("{}…", cut.trim_end())
191 - }
192 -}
193 -
194 -fn invalid(msg: impl Into<String>) -> agent_client_protocol::Error {
195 - agent_client_protocol::Error::invalid_params().data(msg.into())
196 -}
197 -
198 -fn unknown_session(id: &SessionId) -> agent_client_protocol::Error {
199 - invalid(format!("unknown session {id}"))
200 -}
201 -
202 -/// A tool message; `failed` is kept in a private field that never reaches the model.
203 -fn tool_message(id: &str, content: &str, failed: bool) -> Value {
204 - let mut m = json!({ "role": "tool", "tool_call_id": id, "content": content });
205 - if failed {
206 - m["x_failed"] = json!(true);
207 - }
208 - m
209 -}
210 -
211 -fn wire_message(m: &Value) -> Value {
212 - match m.as_object() {
213 - Some(o) if o.contains_key("x_failed") => {
214 - let mut o = o.clone();
215 - o.remove("x_failed");
216 - Value::Object(o)
217 - }
218 - _ => m.clone(),
219 - }
220 -}
221 -
222 -fn locations_of(paths: &[PathBuf]) -> Vec<ToolCallLocation> {
223 - paths
224 - .iter()
225 - .map(|p| ToolCallLocation::new(absolutize(p)))
226 - .collect()
227 -}
228 -
229 -/// Text of a stored user message (`content` is a string, or an array of parts).
230 -fn user_text(msg: &Value) -> String {
231 - match msg.get("content") {
232 - Some(Value::String(s)) => s.clone(),
233 - Some(Value::Array(parts)) => parts
234 - .iter()
235 - .filter_map(|p| match p.get("type").and_then(Value::as_str) {
236 - Some("text") => p.get("text").and_then(Value::as_str).map(String::from),
237 - Some("image_url") => Some("[image]".into()),
238 - _ => None,
239 - })
240 - .collect::<Vec<_>>()
241 - .join("\n\n"),
242 - _ => String::new(),
243 - }
244 -}
245 -
246 -#[derive(Clone)]
247 -pub struct SplitFireAgent {
248 - /// Replaced when credentials change: `authenticate`, `logout`, or a key stored by
249 - /// `--setup` in another process.
250 - llm: Arc<Mutex<LlmClient>>,
251 - /// Process default for each session's `auto_approve` (`--yolo` / SPLITFIRE_YOLO).
252 - yolo: bool,
253 - /// How the agent was launched: `tui` (interactive terminal UI) or `acp` (editor).
254 - surface: &'static str,
255 - client_caps: Arc<Mutex<ClientCapabilities>>,
256 - sessions: Arc<Mutex<HashMap<SessionId, Session>>>,
257 - /// Models offered for selection, fetched from `GET /models` on first use and dropped
258 - /// when credentials change.
259 - models: Arc<tokio::sync::Mutex<Option<Vec<String>>>>,
260 - store: SessionStore,
261 - data_dir: PathBuf,
262 -}
263 -
264 -impl SplitFireAgent {
265 - pub fn new(config: LlmConfig, yolo: bool, surface: &'static str, data_dir: PathBuf) -> Self {
266 - Self {
267 - llm: Arc::new(Mutex::new(LlmClient::new(config))),
268 - yolo,
269 - surface,
270 - client_caps: Arc::default(),
271 - sessions: Arc::default(),
272 - models: Arc::default(),
273 - store: SessionStore::new(&data_dir),
274 - data_dir,
275 - }
276 - }
277 -
278 - pub fn set_client_caps(&self, caps: ClientCapabilities) {
279 - *self.client_caps.lock().unwrap() = caps;
280 - }
281 -
282 - fn llm(&self) -> LlmClient {
283 - self.llm.lock().unwrap().clone()
284 - }
285 -
286 - pub fn has_api_key(&self) -> bool {
287 - self.llm().has_api_key()
288 - }
289 -
290 - pub fn model(&self) -> String {
291 - self.llm().model().to_string()
292 - }
293 -
294 - /// Re-read credentials (environment + stored `env` file) and report whether a key is
295 - /// configured. Editors run terminal auth (`splitfire-agent --setup`) in a separate process
296 - /// and then retry `session/new` or `session/load` on this one without calling
297 - /// `authenticate`, so the stored key has to be picked up here. The model list is only
298 - /// dropped when the configuration actually changed.
299 - pub async fn refresh_credentials(&self) -> bool {
300 - let fresh = LlmConfig::from_env();
301 - let changed = fresh != *self.llm().config();
302 - if changed {
303 - tracing::info!("credentials changed; reloading the Onde client");
304 - *self.llm.lock().unwrap() = LlmClient::new(fresh);
305 - *self.models.lock().await = None;
306 - }
307 - self.has_api_key()
308 - }
309 -
310 - /// `authenticate`: the client has run `splitfire-agent --setup`; pick up the stored key.
311 - pub async fn authenticate(
312 - &self,
313 - req: AuthenticateRequest,
314 - ) -> agent_client_protocol::Result<AuthenticateResponse> {
315 - if &*req.method_id.0 != AUTH_METHOD_ID {
316 - return Err(invalid(format!("unknown auth method {}", req.method_id.0)));
317 - }
318 - if self.refresh_credentials().await {
319 - Ok(AuthenticateResponse::new())
320 - } else {
321 - Err(auth_required_error())
322 - }
323 - }
324 -
325 - /// `logout`: remove the stored key and forget it. A key exported in the agent's own
326 - /// environment can't be removed from here and keeps the agent signed in.
327 - pub async fn logout(&self) -> agent_client_protocol::Result<LogoutResponse> {
328 - let removed = match crate::llm::config_file() {
329 - Some(path) => match std::fs::remove_file(&path) {
330 - Err(e) if e.kind() != std::io::ErrorKind::NotFound => {
331 - Err(format!("removing {}: {e}", path.display()))
332 - }
333 - _ => Ok(()),
334 - },
335 - None => Ok(()),
336 - };
337 - self.refresh_credentials().await;
338 - removed.map_err(|e| agent_client_protocol::Error::internal_error().data(e))?;
339 - Ok(LogoutResponse::new())
340 - }
341 -
342 - fn scratch_dir(&self, id: &SessionId) -> Option<PathBuf> {
343 - let id = crate::session_store::sanitize_id(&id.0).ok()?;
344 - Some(self.data_dir.join("scratch").join(id))
345 - }
346 -
347 - fn remove_scratch(&self, id: &SessionId) {
348 - if let Some(dir) = self.scratch_dir(id) {
349 - let _ = std::fs::remove_dir_all(dir);
350 - }
351 - }
352 -
353 - fn system_prompt(&self, cwd: &Path, roots: &[PathBuf], stems: bool) -> String {
354 - prompt::system_prompt(self.surface, cwd, roots, stems)
355 - }
356 -
357 - /// The selectable models. Falls back to just the configured model if the endpoint
358 - /// can't list them; the configured model is always offered.
359 - async fn models(&self) -> Vec<String> {
360 - let mut cache = self.models.lock().await;
361 - if let Some(ids) = &*cache {
362 - return ids.clone();
363 - }
364 - let llm = self.llm();
365 - let default = llm.model().to_string();
366 - // An explicit list wins; editors show this list as the model picker.
367 - let configured: Vec<String> = std::env::var("SPLITFIRE_MODELS")
368 - .unwrap_or_default()
369 - .split(',')
370 - .map(|m| m.trim().to_string())
371 - .filter(|m| !m.is_empty())
372 - .collect();
373 - let mut ids = if !configured.is_empty() {
374 - configured
375 - } else {
376 - match llm.models().await {
377 - Ok(models) => models.into_iter().map(|m| m.id).collect(),
378 - Err(e) => {
379 - tracing::warn!("listing models failed, offering only {default}: {e:#}");
380 - Vec::new()
381 - }
382 - }
383 - };
384 - if !ids.contains(&default) {
385 - ids.insert(0, default);
386 - }
387 - *cache = Some(ids.clone());
388 - ids
389 - }
390 -
391 - /// Whether the client advertised `session.configOptions.boolean`, which ACP requires
392 - /// before an agent may offer `type: "boolean"` options.
393 - fn boolean_options_supported(&self) -> bool {
394 - self.client_caps
395 - .lock()
396 - .unwrap()
397 - .session
398 - .as_ref()
399 - .and_then(|s| s.config_options.as_ref())
400 - .is_some_and(|c| c.boolean.is_some())
401 - }
402 -
403 - async fn config_options(&self, model: &str, auto_approve: bool) -> Vec<SessionConfigOption> {
404 - let options = self
405 - .models()
406 - .await
407 - .into_iter()
408 - .map(|id| SessionConfigSelectOption::new(id.clone(), id))
409 - .collect();
410 - let mut config = vec![
411 - SessionConfigOption::select(
412 - MODEL_CONFIG_ID,
413 - "Model",
414 - model.to_string(),
415 - SessionConfigSelectOptions::Ungrouped(options),
416 - )
417 - .category(SessionConfigOptionCategory::Model),
418 - ];
419 - if self.boolean_options_supported() {
420 - config.push(
421 - SessionConfigOption::boolean(
422 - AUTO_APPROVE_CONFIG_ID,
423 - "Auto-approve actions",
424 - auto_approve,
425 - )
426 - .description(
427 - "Edit files, run commands and separate stems without asking for permission",
428 - ),
429 - );
430 - }
431 - config
432 - }
433 -
434 - /// The model and auto-approve setting of a live session.
435 - fn session_settings(&self, id: &SessionId) -> Option<(String, bool)> {
436 - self.sessions
437 - .lock()
438 - .unwrap()
439 - .get(id)
440 - .map(|s| (s.model.clone(), s.auto_approve.load(Ordering::Relaxed)))
441 - }
442 -
443 - async fn session_config_options(
444 - &self,
445 - id: &SessionId,
446 - ) -> agent_client_protocol::Result<Vec<SessionConfigOption>> {
447 - let (model, auto_approve) = self
448 - .session_settings(id)
449 - .ok_or_else(|| unknown_session(id))?;
450 - Ok(self.config_options(&model, auto_approve).await)
451 - }
452 -
453 - pub async fn set_config_option(
454 - &self,
455 - req: SetSessionConfigOptionRequest,
456 - connection: &ConnectionTo<Client>,
457 - ) -> agent_client_protocol::Result<SetSessionConfigOptionResponse> {
458 - let id = req.session_id.clone();
459 - self.restore(&id).await?;
460 - match &*req.config_id.0 {
461 - MODEL_CONFIG_ID => {
462 - let Some(model) = req.value.as_value_id().map(|v| v.0.to_string()) else {
463 - return Err(invalid("model must be a value id"));
464 - };
465 - if !self.models().await.contains(&model) {
466 - return Err(invalid(format!("unknown model {model}")));
467 - }
468 - match self.sessions.lock().unwrap().get_mut(&id) {
469 - Some(s) => s.model = model.clone(),
470 - None => return Err(unknown_session(&id)),
471 - }
472 - self.persist(&id).await;
473 - tracing::info!("session {id} now uses model {model}");
474 - }
475 - AUTO_APPROVE_CONFIG_ID if self.boolean_options_supported() => {
476 - let Some(on) = req.value.as_bool() else {
477 - return Err(invalid("auto_approve must be a boolean"));
478 - };
479 - match self.sessions.lock().unwrap().get(&id) {
480 - Some(s) => s.auto_approve.store(on, Ordering::Relaxed),
481 - None => return Err(unknown_session(&id)),
482 - }
483 - tracing::info!("session {id} auto-approve {on}");
484 - }
485 - other => return Err(invalid(format!("unknown config option {other}"))),
486 - }
487 - let options = self.session_config_options(&id).await?;
488 - // Tell every view of the session about the change, not just the requester.
489 - connection.send_notification(SessionNotification::new(
490 - id,
491 - SessionUpdate::ConfigOptionUpdate(ConfigOptionUpdate::new(options.clone())),
492 - ))?;
493 - Ok(SetSessionConfigOptionResponse::new(options))
494 - }
495 -
496 - /// Save a session that is not running a turn (a turn checks the history out; the
497 - /// end-of-turn save covers it).
498 - async fn persist(&self, id: &SessionId) {
499 - let stored = self
500 - .sessions
501 - .lock()
502 - .unwrap()
503 - .get(id)
504 - .filter(|s| !s.busy)
505 - .map(Session::stored);
506 - let Some(stored) = stored else { return };
507 - let store = self.store.clone();
508 - let sid = id.0.to_string();
509 - match tokio::task::spawn_blocking(move || store.save(&sid, &stored)).await {
510 - Ok(Ok(())) => {}
511 - Ok(Err(e)) => tracing::warn!("saving session {id} failed: {e:#}"),
512 - Err(e) => tracing::warn!("saving session {id} failed: {e}"),
513 - }
514 - }
515 -
516 - /// Bring a stored session into memory without rebinding it (a prompt or config change
517 - /// for a thread the client reopened without `session/load`). Errors if it is unknown.
518 - async fn restore(&self, id: &SessionId) -> agent_client_protocol::Result<()> {
519 - if self.sessions.lock().unwrap().contains_key(id) {
520 - return Ok(());
521 - }
522 - let stored = self
523 - .load_stored(id)
524 - .await?
525 - .ok_or_else(|| unknown_session(id))?;
526 - tracing::info!("restored session {id} from disk");
527 - self.sessions
528 - .lock()
529 - .unwrap()
530 - .entry(id.clone())
531 - .or_insert_with(|| Session::restored(stored, self.yolo));
532 - Ok(())
533 - }
534 -
535 - async fn load_stored(
536 - &self,
537 - id: &SessionId,
538 - ) -> agent_client_protocol::Result<Option<StoredSession>> {
539 - let store = self.store.clone();
540 - let sid = id.0.to_string();
541 - tokio::task::spawn_blocking(move || store.load(&sid))
542 - .await
543 - .map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))?
544 - .map_err(|e| agent_client_protocol::Error::internal_error().data(format!("{e:#}")))
545 - }
546 -
547 - fn stems_available(&self, id: &SessionId) -> bool {
548 - self.sessions
549 - .lock()
550 - .unwrap()
551 - .get(id)
552 - .is_some_and(|s| s.mcp.has_stem_separation())
553 - }
554 -
555 - /// `AvailableCommandsUpdate` for a session, sent after the response that created it.
556 - pub fn send_commands(
557 - &self,
558 - connection: &ConnectionTo<Client>,
559 - id: &SessionId,
560 - ) -> agent_client_protocol::Result<()> {
561 - let commands = prompt::slash_commands(self.stems_available(id));
562 - connection.send_notification(SessionNotification::new(
563 - id.clone(),
564 - SessionUpdate::AvailableCommandsUpdate(AvailableCommandsUpdate::new(commands)),
565 - ))
566 - }
567 -
568 - pub async fn new_session(
569 - &self,
570 - req: NewSessionRequest,
571 - ) -> agent_client_protocol::Result<NewSessionResponse> {
572 - check_absolute(&req.cwd, &req.additional_directories)?;
573 - // Lexical normalization only: the client sees `cwd` again in `session/list`, so
574 - // symlinks in it are not rewritten.
575 - let cwd = normalize(&req.cwd);
576 - let roots: Vec<PathBuf> = req
577 - .additional_directories
578 - .iter()
579 - .map(|r| normalize(r))
580 - .collect();
581 - let id = SessionId::new(uuid::Uuid::new_v4().to_string());
582 - let model = self.model();
583 - let mut session = Session::fresh(cwd, roots, model.clone(), self.yolo);
584 - session.mcp = Arc::new(McpToolset::connect(&req.mcp_servers).await);
585 - self.sessions.lock().unwrap().insert(id.clone(), session);
586 - // Saved right away so the thread survives a restart before its first prompt.
587 - self.persist(&id).await;
588 - let options = self.config_options(&model, self.yolo).await;
589 - Ok(NewSessionResponse::new(id).config_options(options))
590 - }
591 -
592 - /// Live and stored sessions, optionally filtered by working directory, newest first.
593 - pub fn list_sessions(
594 - &self,
595 - req: ListSessionsRequest,
596 - ) -> agent_client_protocol::Result<ListSessionsResponse> {
597 - // Every session fits on one page, so no cursor is ever handed out; any cursor is stale.
598 - if let Some(cursor) = &req.cursor {
599 - return Err(invalid(format!("invalid cursor {cursor}")));
600 - }
601 - struct Row {
602 - id: SessionId,
603 - cwd: PathBuf,
604 - roots: Vec<PathBuf>,
605 - title: Option<String>,
606 - updated_at: u64,
607 - }
608 - let mut rows: HashMap<String, Row> = self
609 - .store
610 - .list()
611 - .into_iter()
612 - .map(|(id, s)| {
613 - (
614 - id.clone(),
615 - Row {
616 - id: SessionId::new(id),
617 - cwd: s.cwd,
618 - roots: s.roots,
619 - title: s.title,
620 - updated_at: s.updated_at,
621 - },
622 - )
623 - })
624 - .collect();
625 - // Live state wins over the (possibly older) stored copy.
626 - for (id, s) in self.sessions.lock().unwrap().iter() {
627 - rows.insert(
628 - id.0.to_string(),
629 - Row {
630 - id: id.clone(),
631 - cwd: s.cwd.clone(),
632 - roots: s.roots.clone(),
633 - title: s.title.clone(),
634 - updated_at: s.updated_at,
635 - },
636 - );
637 - }
638 - // Stored cwds are normalized, so normalize the filter the same way.
639 - let wanted = req.cwd.as_deref().map(normalize);
640 - let mut rows: Vec<Row> = rows
641 - .into_values()
642 - .filter(|r| wanted.as_ref().is_none_or(|cwd| *cwd == normalize(&r.cwd)))
643 - .collect();
644 - rows.sort_by(|a, b| {
645 - b.updated_at
646 - .cmp(&a.updated_at)
647 - .then_with(|| a.id.0.cmp(&b.id.0))
648 - });
649 - rows.truncate(MAX_LISTED);
650 - Ok(ListSessionsResponse::new(
651 - rows.into_iter()
652 - .map(|r| {
653 - SessionInfo::new(r.id, r.cwd)
654 - .additional_directories(r.roots)
655 - .title(r.title)
656 - .updated_at(iso8601(r.updated_at))
657 - })
658 - .collect(),
659 - ))
660 - }
661 -
662 - pub fn cancel(&self, session_id: &SessionId) {
663 - if let Some(s) = self.sessions.lock().unwrap().get(session_id) {
664 - s.cancel.cancel();
665 - }
666 - }
667 -
668 - /// Make `id` live for `session/load` and `session/resume`: from memory, else from the
669 - /// store; rebind cwd and roots, and reconnect MCP servers from the request. A well-formed
670 - /// id with no saved history (a thread from a process that never saved it) reopens empty
671 - /// under the same id, so the editor's thread opens instead of failing to launch. Returns
672 - /// whether the history was lost.
673 - async fn reopen(
674 - &self,
675 - id: &SessionId,
676 - cwd: &Path,
677 - roots: &[PathBuf],
678 - mcp_servers: &[McpServer],
679 - ) -> agent_client_protocol::Result<bool> {
680 - check_absolute(cwd, roots)?;
681 - if crate::session_store::sanitize_id(&id.0).is_err() {
682 - return Err(unknown_session(id));
683 - }
684 - let live = self.sessions.lock().unwrap().contains_key(id);
685 - let stored = if live {
686 - None
687 - } else {
688 - self.load_stored(id).await?
689 - };
690 - let history_lost = !live && stored.is_none();
691 - let cwd = normalize(cwd);
692 - let roots: Vec<PathBuf> = roots.iter().map(|r| normalize(r)).collect();
693 - let mcp = Arc::new(McpToolset::connect(mcp_servers).await);
694 -
695 - let old_mcp = {
696 - let mut sessions = self.sessions.lock().unwrap();
697 - if !sessions.contains_key(id) {
698 - let session = match stored {
699 - Some(stored) => Session::restored(stored, self.yolo),
700 - None => {
701 - tracing::warn!("session {id} has no saved history; reopening it empty");
702 - Session::fresh(cwd.clone(), roots.clone(), self.model(), self.yolo)
703 - }
704 - };
705 - sessions.insert(id.clone(), session);
706 - }
707 - let s = sessions.get_mut(id).ok_or_else(|| unknown_session(id))?;
708 - s.cwd = cwd;
709 - if !roots.is_empty() {
710 - s.roots = roots;
711 - }
712 - s.updated_at = now_secs();
713 - std::mem::replace(&mut s.mcp, mcp)
714 - };
715 - if let Some(old) = Arc::into_inner(old_mcp) {
716 - tokio::spawn(old.shutdown());
717 - }
718 - self.persist(id).await;
719 - Ok(history_lost)
720 - }
721 -
722 - /// Tell the user a reopened thread starts without its earlier history. Shown in the
723 - /// thread only; the model doesn't see it.
724 - fn notify_history_lost(
725 - &self,
726 - connection: &ConnectionTo<Client>,
727 - id: &SessionId,
728 - ) -> agent_client_protocol::Result<()> {
729 - connection.send_notification(SessionNotification::new(
730 - id.clone(),
731 - SessionUpdate::AgentMessageChunk(
732 - ContentChunk::new(
733 - "SplitFire no longer has this conversation's history, so it starts fresh \
734 - from here."
735 - .to_string()
736 - .into(),
737 - )
738 - .message_id(new_message_id()),
739 - ),
740 - ))
741 - }
742 -
743 - /// Replay a session's conversation as `session/update` notifications (text, tool calls
744 - /// with their status and output), then respond with config options. `session/load`.
745 - pub async fn load_session(
746 - &self,
747 - req: LoadSessionRequest,
748 - connection: &ConnectionTo<Client>,
749 - ) -> agent_client_protocol::Result<LoadSessionResponse> {
750 - let id = req.session_id.clone();
751 - let history_lost = self
752 - .reopen(&id, &req.cwd, &req.additional_directories, &req.mcp_servers)
753 - .await?;
754 - let (messages, cwd, roots) = {
755 - let sessions = self.sessions.lock().unwrap();
756 - let s = sessions.get(&id).ok_or_else(|| unknown_session(&id))?;
757 - (s.messages.clone(), s.cwd.clone(), s.roots.clone())
758 - };
759 - let notify = |update: SessionUpdate| {
760 - connection.send_notification(SessionNotification::new(id.clone(), update))
761 - };
762 - // Tool output by call id, from the stored `tool` messages.
763 - let results: HashMap<&str, (&str, bool)> = messages
764 - .iter()
765 - .filter(|m| m.get("role").and_then(Value::as_str) == Some("tool"))
766 - .filter_map(|m| {
767 - Some((
768 - m.get("tool_call_id")?.as_str()?,
769 - (
770 - m.get("content").and_then(Value::as_str).unwrap_or(""),
771 - m.get("x_failed").and_then(Value::as_bool).unwrap_or(false),
772 - ),
773 - ))
774 - })
775 - .collect();
776 - for msg in &messages {
777 - match msg.get("role").and_then(Value::as_str).unwrap_or("") {
778 - "user" => {
779 - let text = user_text(msg);
780 - if !text.is_empty() {
781 - notify(SessionUpdate::UserMessageChunk(
782 - ContentChunk::new(text.into()).message_id(new_message_id()),
783 - ))?;
784 - }
785 - }
786 - "assistant" => {
787 - let text = msg.get("content").and_then(Value::as_str).unwrap_or("");
788 - if !text.is_empty() {
789 - notify(SessionUpdate::AgentMessageChunk(
790 - ContentChunk::new(text.to_string().into()).message_id(new_message_id()),
791 - ))?;
792 - }
793 - let calls = msg.get("tool_calls").and_then(Value::as_array);
794 - for tc in calls.into_iter().flatten() {
795 - let call_id = tc.get("id").and_then(Value::as_str).unwrap_or("unknown");
796 - let name = tc
797 - .pointer("/function/name")
798 - .and_then(Value::as_str)
799 - .unwrap_or("unknown");
800 - let args: Value = tc
801 - .pointer("/function/arguments")
802 - .and_then(Value::as_str)
803 - .and_then(|a| serde_json::from_str(a).ok())
804 - .unwrap_or_else(|| json!({}));
805 - let (title, kind, locations) = tools::describe(&cwd, &roots, name, &args);
806 - let mut call = ToolCall::new(call_id.to_string(), title)
807 - .name(name.to_string())
808 - .kind(kind)
809 - .locations(locations)
810 - .raw_input(args);
811 - match results.get(call_id) {
812 - Some((output, failed)) => {
813 - call = call
814 - .status(if *failed {
815 - ToolCallStatus::Failed
816 - } else {
817 - ToolCallStatus::Completed
818 - })
819 - .content(vec![ToolCallContent::from(output.to_string())])
820 - .raw_output(json!({ "output": output, "failed": failed }));
821 - }
822 - // The turn was interrupted before this call produced a result.
823 - None => call = call.status(ToolCallStatus::Failed),
824 - }
825 - notify(SessionUpdate::ToolCall(call))?;
826 - }
827 - }
828 - // system and tool messages are not replayed on their own.
829 - _ => {}
830 - }
831 - }
832 - if history_lost {
833 - self.notify_history_lost(connection, &id)?;
834 - }
835 - Ok(LoadSessionResponse::new().config_options(self.session_config_options(&id).await?))
836 - }
837 -
838 - /// Rebind a session to a new cwd / additional directories without replaying it.
839 - pub async fn resume_session(
840 - &self,
841 - req: ResumeSessionRequest,
842 - connection: &ConnectionTo<Client>,
843 - ) -> agent_client_protocol::Result<ResumeSessionResponse> {
844 - let id = req.session_id.clone();
845 - let history_lost = self
846 - .reopen(&id, &req.cwd, &req.additional_directories, &req.mcp_servers)
847 - .await?;
848 - if history_lost {
849 - self.notify_history_lost(connection, &id)?;
850 - }
851 - Ok(ResumeSessionResponse::new().config_options(self.session_config_options(&id).await?))
852 - }
853 -
854 - /// Cancel any in-progress work, release MCP servers and scratch audio. The stored
855 - /// conversation stays; `session/delete` removes it.
856 - pub fn close_session(
857 - &self,
858 - req: CloseSessionRequest,
859 - ) -> agent_client_protocol::Result<CloseSessionResponse> {
860 - let id = req.session_id.clone();
861 - let removed = self.sessions.lock().unwrap().remove(&id);
862 - match removed {
863 - Some(s) => self.release(&id, s),
864 - // A session from an earlier process that was never reopened has nothing to close.
865 - None if self.store.load(&id.0).ok().flatten().is_some() => {}
866 - None => return Err(unknown_session(&id)),
867 - }
868 - Ok(CloseSessionResponse::new())
869 - }
870 -
871 - /// Forget a session, in memory and on disk, with its scratch audio. Idempotent: deleting
872 - /// an unknown session succeeds.
873 - pub fn delete_session(
874 - &self,
875 - req: DeleteSessionRequest,
876 - ) -> agent_client_protocol::Result<DeleteSessionResponse> {
877 - let id = req.session_id.clone();
878 - if let Some(s) = self.sessions.lock().unwrap().remove(&id) {
879 - self.release(&id, s);
880 - }
881 - // A malformed id can't name a stored session; nothing to remove.
882 - if crate::session_store::sanitize_id(&id.0).is_ok() {
883 - self.store.delete(&id.0).map_err(|e| {
884 - agent_client_protocol::Error::internal_error().data(format!("{e:#}"))
885 - })?;
886 - }
887 - self.remove_scratch(&id);
888 - Ok(DeleteSessionResponse::new())
889 - }
890 -
891 - fn release(&self, id: &SessionId, s: Session) {
892 - s.cancel.cancel();
893 - self.remove_scratch(id);
894 - if let Some(mcp) = Arc::into_inner(s.mcp) {
895 - tokio::spawn(mcp.shutdown());
896 - }
897 - }
898 -
899 - /// Stop everything: used when the connection ends.
900 - pub async fn shutdown(&self) {
901 - let sessions: Vec<Session> = self
902 - .sessions
903 - .lock()
904 - .unwrap()
905 - .drain()
906 - .map(|(_, s)| s)
907 - .collect();
908 - for s in sessions {
909 - s.cancel.cancel();
910 - if let Some(mcp) = Arc::into_inner(s.mcp) {
911 - mcp.shutdown().await;
912 - }
913 - }
914 - mcp::shutdown_demucs().await;
915 - }
916 -
917 - pub async fn prompt(
918 - &self,
919 - request: PromptRequest,
920 - responder: Responder<PromptResponse>,
921 - connection: ConnectionTo<Client>,
922 - ) -> agent_client_protocol::Result<()> {
923 - // A client may prompt a reopened thread without loading or resuming it first.
924 - if let Err(e) = self.restore(&request.session_id).await {
925 - return responder.respond_with_error(e);
926 - }
927 - match self.run_turn(request, connection).await {
928 - Ok(stop) => responder.respond(PromptResponse::new(stop)),
929 - Err(e) if e.downcast_ref::<UnknownSession>().is_some() => {
930 - responder.respond_with_error(invalid(e.to_string()))
931 - }
932 - Err(e) => responder.respond_with_error(
933 - agent_client_protocol::Error::internal_error().data(format!("{e:#}")),
934 - ),
935 - }
936 - }
937 -
938 - async fn run_turn(
939 - &self,
940 - request: PromptRequest,
941 - connection: ConnectionTo<Client>,
942 - ) -> anyhow::Result<StopReason> {
943 - let session_id = request.session_id.clone();
944 - let caps = self.client_caps.lock().unwrap().clone();
945 - let blocks =
946 - prompt::resolve_resource_links(&request.prompt, &caps, &connection, &session_id).await;
947 - let prompt_text = prompt::prompt_to_text(&blocks);
948 - // Check the session out for the turn.
949 - let (
950 - cwd,
951 - roots,
952 - model,
953 - history,
954 - cancel,
955 - always_allowed,
956 - always_rejected,
957 - auto_approve,
958 - mcp,
959 - new_title,
960 - ) = {
961 - let mut sessions = self.sessions.lock().unwrap();
962 - let s = sessions
963 - .get_mut(&session_id)
964 - .ok_or_else(|| UnknownSession(session_id.to_string()))?;
965 - if s.busy {
966 - anyhow::bail!("a prompt is already running for session {session_id}");
967 - }
968 - s.busy = true;
969 - s.cancel = CancellationToken::new();
970 - s.updated_at = now_secs();
971 - let new_title = if s.title.is_none() && !prompt_text.trim().is_empty() {
972 - s.title = Some(session_title(&prompt_text));
973 - s.title.clone()
974 - } else {
975 - None
976 - };
977 - (
978 - s.cwd.clone(),
979 - s.roots.clone(),
980 - s.model.clone(),
981 - std::mem::take(&mut s.messages),
982 - s.cancel.clone(),
983 - s.always_allowed.clone(),
984 - s.always_rejected.clone(),
985 - s.auto_approve.clone(),
986 - s.mcp.clone(),
987 - new_title,
988 - )
989 - };
990 - if let Some(title) = new_title {
991 - let update = SessionInfoUpdate::new()
992 - .title(title)
993 - .updated_at(iso8601(now_secs()));
994 - let _ = connection.send_notification(SessionNotification::new(
995 - session_id.clone(),
996 - SessionUpdate::SessionInfoUpdate(update),
997 - ));
998 - }
999 -
1000 - let stems = mcp.has_stem_separation();
1001 - let mut messages = vec![json!({
1002 - "role": "system",
1003 - "content": self.system_prompt(&cwd, &roots, stems),
1004 - })];
1005 - messages.extend(history);
1006 -
1007 - let ctx = ToolCtx {
1008 - connection: connection.clone(),
1009 - session_id: session_id.clone(),
1010 - cwd,
1011 - roots,
1012 - caps,
1013 - cancel,
1014 - auto_approve,
1015 - always_allowed,
1016 - always_rejected,
1017 - mcp,
1018 - };
1019 - let result = self.turn(&ctx, &model, &mut messages, &blocks, stems).await;
1020 -
1021 - // Check the history back in, even when the turn failed.
1022 - let history = messages.split_off(1);
1023 - let stored = {
1024 - let mut sessions = self.sessions.lock().unwrap();
1025 - sessions.get_mut(&session_id).map(|s| {
1026 - s.messages = history;
1027 - s.busy = false;
1028 - s.updated_at = now_secs();
1029 - s.stored()
1030 - })
1031 - };
1032 - if let Some(stored) = stored {
1033 - let updated_at = stored.updated_at;
1034 - let store = self.store.clone();
1035 - let id = session_id.0.to_string();
1036 - let saved = tokio::task::spawn_blocking(move || store.save(&id, &stored)).await;
1037 - if let Ok(Err(e)) | Err(e) = saved.map_err(anyhow::Error::from) {
1038 - tracing::warn!("saving session {session_id} failed: {e:#}");
1039 - }
1040 - let _ = connection.send_notification(SessionNotification::new(
1041 - session_id.clone(),
1042 - SessionUpdate::SessionInfoUpdate(
1043 - SessionInfoUpdate::new().updated_at(iso8601(updated_at)),
1044 - ),
1045 - ));
1046 - }
1047 - result
1048 - }
1049 -
1050 - /// Everything inside a turn after the session is checked out.
1051 - async fn turn(
1052 - &self,
1053 - ctx: &ToolCtx,
1054 - model: &str,
1055 - messages: &mut Vec<Value>,
1056 - blocks: &[ContentBlock],
1057 - stems: bool,
1058 - ) -> anyhow::Result<StopReason> {
1059 - let audio_paths = match self.scratch_dir(&ctx.session_id) {
1060 - Some(dir) => prompt::save_audio_blocks(blocks, &dir).await?,
1061 - None => Vec::new(),
1062 - };
1063 - let slash = prompt::expand_slash(&prompt::prompt_to_text(blocks), stems);
1064 - let content = prompt::prompt_to_content(blocks, &audio_paths, slash.as_deref());
1065 - messages.push(json!({ "role": "user", "content": content }));
1066 - self.agent_loop(ctx, model, messages).await
1067 - }
1068 -
1069 - async fn agent_loop(
1070 - &self,
1071 - ctx: &ToolCtx,
1072 - model: &str,
1073 - messages: &mut Vec<Value>,
1074 - ) -> anyhow::Result<StopReason> {
1075 - let llm = self.llm();
1076 - let tool_defs = tools::definitions(&ctx.mcp);
1077 - let notify = |update: SessionUpdate| {
1078 - ctx.connection
1079 - .send_notification(SessionNotification::new(ctx.session_id.clone(), update))
1080 - };
1081 - // `size` is the model's context window, which OpenAI-style endpoints don't report.
1082 - let context_window = std::env::var("SPLITFIRE_CONTEXT_WINDOW")
1083 - .ok()
1084 - .and_then(|v| v.parse().ok())
1085 - .unwrap_or(DEFAULT_CONTEXT_WINDOW);
1086 -
1087 - for _ in 0..MAX_TURNS {
1088 - // One message id per assistant message (and per thought), so clients can group
1089 - // the streamed chunks.
1090 - let message_id = new_message_id();
1091 - let thought_id = new_message_id();
1092 - let wire: Vec<Value> = messages.iter().map(wire_message).collect();
1093 - let completion = tokio::select! {
1094 - r = llm.complete(model, &wire, &tool_defs, |delta| {
1095 - let update = match delta {
1096 - Delta::Text(t) => SessionUpdate::AgentMessageChunk(
1097 - ContentChunk::new(t.to_string().into()).message_id(message_id.clone()),
1098 - ),
1099 - Delta::Reasoning(t) => SessionUpdate::AgentThoughtChunk(
1100 - ContentChunk::new(t.to_string().into()).message_id(thought_id.clone()),
1101 - ),
1102 - };
1103 - if let Err(e) = notify(update) {
1104 - tracing::warn!("failed to send update: {e}");
1105 - }
1106 - }) => r?,
1107 - () = ctx.cancel.cancelled() => return Ok(StopReason::Cancelled),
1108 - };
1109 - messages.push(completion.to_message());
1110 - if let Some(usage) = completion.usage {
1111 - let used = usage.used();
1112 - notify(SessionUpdate::UsageUpdate(UsageUpdate::new(
1113 - used,
1114 - context_window.max(used),
1115 - )))?;
1116 - }
1117 -
1118 - if completion.tool_calls.is_empty() {
1119 - return Ok(match completion.finish_reason.as_deref() {
1120 - Some("length") => StopReason::MaxTokens,
1121 - Some("content_filter") => StopReason::Refusal,
1122 - _ => StopReason::EndTurn,
1123 - });
1124 - }
1125 -
1126 - for tc in &completion.tool_calls {
1127 - // Every tool call needs a matching tool message, even after cancellation.
1128 - if ctx.cancel.is_cancelled() {
1129 - messages.push(tool_message(&tc.id, "Cancelled by user.", true));
1130 - continue;
1131 - }
1132 - let parsed: Result<Value, _> = serde_json::from_str(&tc.arguments);
1133 - let args = parsed.as_ref().cloned().unwrap_or_else(|_| json!({}));
1134 - let (title, kind, locations) =
1135 - tools::describe(&ctx.cwd, &ctx.roots, &tc.name, &args);
1136 - notify(SessionUpdate::ToolCall(
1137 - ToolCall::new(tc.id.clone(), title)
1138 - .name(tc.name.clone())
1139 - .kind(kind)
1140 - .status(ToolCallStatus::InProgress)
1141 - .locations(locations)
1142 - .raw_input(args.clone()),
1143 - ))?;
1144 -
1145 - let outcome: ToolOutcome = if parsed.is_err() {
1146 - ToolOutcome::err(format!("Invalid JSON arguments: {}", tc.arguments))
1147 - } else {
1148 - tools::execute(ctx, &tc.id, &tc.name, args).await
1149 - };
1150 -
1151 - let mut fields = ToolCallUpdateFields::new()
1152 - .status(if outcome.failed {
1153 - ToolCallStatus::Failed
1154 - } else {
1155 - ToolCallStatus::Completed
1156 - })
1157 - .content(outcome.content)
1158 - .raw_output(json!({ "output": outcome.text, "failed": outcome.failed }));
1159 - if !outcome.locations.is_empty() {
1160 - fields = fields.locations(locations_of(&outcome.locations));
1161 - }
1162 - ctx.update(&tc.id, fields)?;
1163 - messages.push(tool_message(&tc.id, &outcome.text, outcome.failed));
1164 - }
1165 - if ctx.cancel.is_cancelled() {
1166 - return Ok(StopReason::Cancelled);
1167 - }
1168 - }
1169 - Ok(StopReason::MaxTurnRequests)
1170 - }
1171 -}
1172 -
1173 -/// Marker error: the prompt named a session that doesn't exist.
1174 -#[derive(Debug)]
1175 -struct UnknownSession(String);
1176 -
1177 -impl std::fmt::Display for UnknownSession {
1178 - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1179 - write!(f, "unknown session {}", self.0)
1180 - }
1181 -}
1182 -
1183 -impl std::error::Error for UnknownSession {}
1184 -
1185 -#[cfg(test)]
1186 -mod tests {
1187 - use super::*;
1188 -
1189 - #[test]
1190 - fn iso8601_known_dates() {
1191 - assert_eq!(iso8601(0), "1970-01-01T00:00:00Z");
1192 - assert_eq!(iso8601(1_700_000_000), "2023-11-14T22:13:20Z");
1193 - assert_eq!(iso8601(951_782_400), "2000-02-29T00:00:00Z");
1194 - }
1195 -
1196 - #[test]
1197 - fn titles_are_trimmed_and_single_line() {
1198 - assert_eq!(session_title(" hello\nworld"), "hello");
1199 - let long = "x".repeat(200);
1200 - let t = session_title(&long);
1201 - assert_eq!(t.chars().count(), 80);
1202 - assert!(t.ends_with('…'));
1203 - }
1204 -
1205 - #[test]
1206 - fn version_negotiation_always_v1() {
1207 - assert_eq!(negotiate_version(ProtocolVersion::V1), ProtocolVersion::V1);
1208 - assert_eq!(
1209 - negotiate_version(ProtocolVersion::from(7)),
1210 - ProtocolVersion::V1
1211 - );
1212 - assert_eq!(
1213 - negotiate_version(ProtocolVersion::from(0)),
1214 - ProtocolVersion::V1
1215 - );
1216 - }
1217 -
1218 - #[test]
1219 - fn private_failure_flag_never_reaches_the_wire() {
1220 - let m = tool_message("c1", "boom", true);
1221 - assert_eq!(m["x_failed"], true);
1222 - let w = wire_message(&m);
1223 - assert!(w.get("x_failed").is_none());
1224 - assert_eq!(w["content"], "boom");
1225 - assert!(tool_message("c1", "ok", false).get("x_failed").is_none());
1226 - }
1227 -
1228 - #[test]
1229 - fn user_text_handles_parts() {
1230 - let m = json!({"role": "user", "content": [
1231 - {"type": "text", "text": "hi"},
1232 - {"type": "image_url", "image_url": {"url": "x"}}
1233 - ]});
1234 - assert_eq!(user_text(&m), "hi\n\n[image]");
1235 - assert_eq!(user_text(&json!({"content": "plain"})), "plain");
1236 - }
1237 -}
src/bin/mcp_stdio_stub.rs
+3 -177
@@ -1,181 +1,7 @@
1 -//! Test-only MCP server over stdio (newline-delimited JSON-RPC). Stands in for third-party
2 -//! servers and for `demucs --mcp`: it ignores its arguments.
3 -//!
4 -//! Tools: `echo` (read-only), `slow_progress` (three progress notifications),
5 -//! `wait_for_cancel` (answers only after `notifications/cancelled`), and the demucs-shaped
6 -//! `list_models` / `separate_stems`.
7 -
8 -use std::collections::HashSet;
9 -use std::sync::{Arc, Mutex};
10 -use std::time::Duration;
11 -
12 -use serde_json::{Value, json};
13 -use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
14 -
15 -fn tools() -> Value {
16 - json!([
17 - {
18 - "name": "echo",
19 - "description": "Echo the text back.",
20 - "inputSchema": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]},
21 - "annotations": {"readOnlyHint": true}
22 - },
23 - {
24 - "name": "slow_progress",
25 - "description": "Report progress three times, then finish.",
26 - "inputSchema": {"type": "object", "properties": {}}
27 - },
28 - {
29 - "name": "wait_for_cancel",
30 - "description": "Blocks until the call is cancelled.",
31 - "inputSchema": {"type": "object", "properties": {}}
32 - },
33 - {
34 - "name": "list_models",
35 - "description": "List stem separation models.",
36 - "inputSchema": {"type": "object", "properties": {}},
37 - "annotations": {"readOnlyHint": true}
38 - },
39 - {
40 - "name": "separate_stems",
41 - "description": "Pretend to separate stems.",
42 - "inputSchema": {"type": "object", "properties": {"input": {"type": "string"}, "model": {"type": "string"}}, "required": ["input"]}
43 - }
44 - ])
45 -}
1 +//! Test MCP server for the integration tests: `ed_mcp::test_server` as a binary, so tests can
2 +//! hand its path to the agent as a client-supplied server or as `SPLITFIRE_DEMUCS_BIN`.
3
4 #[tokio::main]
5 async fn main() {
49 - // Lets tests find out whether this process is still alive after the agent closes it.
50 - if let Some(path) = std::env::var_os("STUB_PID_FILE") {
51 - let _ = std::fs::write(path, std::process::id().to_string());
52 - }
53 - let out = Arc::new(tokio::sync::Mutex::new(tokio::io::stdout()));
54 - let cancelled: Arc<Mutex<HashSet<String>>> = Arc::default();
55 - let mut lines = BufReader::new(tokio::io::stdin()).lines();
56 -
57 - let send = {
58 - let out = out.clone();
59 - move |msg: Value| {
60 - let out = out.clone();
61 - async move {
62 - let mut o = out.lock().await;
63 - let _ = o.write_all(format!("{msg}\n").as_bytes()).await;
64 - let _ = o.flush().await;
65 - }
66 - }
67 - };
68 -
69 - while let Ok(Some(line)) = lines.next_line().await {
70 - let Ok(msg) = serde_json::from_str::<Value>(&line) else {
71 - continue;
72 - };
73 - let method = msg
74 - .get("method")
75 - .and_then(Value::as_str)
76 - .unwrap_or("")
77 - .to_string();
78 - let id = msg.get("id").cloned();
79 - let params = msg.get("params").cloned().unwrap_or(Value::Null);
80 - let Some(id) = id else {
81 - if method == "notifications/cancelled"
82 - && let Some(rid) = params.get("requestId")
83 - {
84 - cancelled.lock().unwrap().insert(rid.to_string());
85 - }
86 - continue;
87 - };
88 - match method.as_str() {
89 - "initialize" => {
90 - let version = params
91 - .get("protocolVersion")
92 - .cloned()
93 - .unwrap_or(json!("2025-06-18"));
94 - send(json!({"jsonrpc": "2.0", "id": id, "result": {
95 - "protocolVersion": version,
96 - "capabilities": {"tools": {}},
97 - "serverInfo": {"name": "stub", "version": "0.0.0"}
98 - }}))
99 - .await;
100 - }
101 - "ping" => send(json!({"jsonrpc": "2.0", "id": id, "result": {}})).await,
102 - "tools/list" => {
103 - send(json!({"jsonrpc": "2.0", "id": id, "result": {"tools": tools()}})).await
104 - }
105 - "tools/call" => {
106 - let send = send.clone();
107 - let cancelled = cancelled.clone();
108 - tokio::spawn(async move {
109 - let name = params
110 - .get("name")
111 - .and_then(Value::as_str)
112 - .unwrap_or("")
113 - .to_string();
114 - let args = params.get("arguments").cloned().unwrap_or(json!({}));
115 - let token = params.pointer("/_meta/progressToken").cloned();
116 - let progress = |n: u32, msg: &str| {
117 - token.as_ref().map(|t| {
118 - json!({"jsonrpc": "2.0", "method": "notifications/progress",
119 - "params": {"progressToken": t, "progress": n, "total": 100, "message": msg}})
120 - })
121 - };
122 - let result = match name.as_str() {
123 - "echo" => {
124 - let text = args.get("text").and_then(Value::as_str).unwrap_or("");
125 - json!({"content": [{"type": "text", "text": format!("echo: {text}")}], "isError": false})
126 - }
127 - "slow_progress" | "separate_stems" => {
128 - for (n, m) in [(10, "starting"), (50, "halfway"), (90, "almost")] {
129 - if let Some(p) = progress(n, m) {
130 - send(p).await;
131 - }
132 - tokio::time::sleep(Duration::from_millis(40)).await;
133 - }
134 - if name == "slow_progress" {
135 - json!({"content": [{"type": "text", "text": "done"}], "isError": false})
136 - } else {
137 - let input = args
138 - .get("input")
139 - .and_then(Value::as_str)
140 - .unwrap_or("/tmp/x.wav");
141 - let dir = format!("{}_stems", input.trim_end_matches(".wav"));
142 - let stems: Vec<Value> = ["drums", "bass", "vocals", "other"]
143 - .iter()
144 - .map(|s| json!({"id": s, "path": format!("{dir}/{s}.wav")}))
145 - .collect();
146 - json!({
147 - "content": [{"type": "text", "text": format!("Separated into {} stems", stems.len())}],
148 - "structuredContent": {"stems": stems},
149 - "isError": false
150 - })
151 - }
152 - }
153 - "list_models" => json!({
154 - "content": [{"type": "text", "text": "htdemucs 80 MB"}],
155 - "structuredContent": {"models": [
156 - {"id": "htdemucs", "size_mb": 80, "cached": false},
157 - {"id": "htdemucs_6s", "size_mb": 55, "cached": true}
158 - ]},
159 - "isError": false
160 - }),
161 - "wait_for_cancel" => {
162 - let key = id.to_string();
163 - while !cancelled.lock().unwrap().contains(&key) {
164 - tokio::time::sleep(Duration::from_millis(20)).await;
165 - }
166 - json!({"content": [{"type": "text", "text": "cancelled"}], "isError": true})
167 - }
168 - other => {
169 - send(json!({"jsonrpc": "2.0", "id": id, "error": {"code": -32602, "message": format!("unknown tool {other}")}})).await;
170 - return;
171 - }
172 - };
173 - send(json!({"jsonrpc": "2.0", "id": id, "result": result})).await;
174 - });
175 - }
176 - other => {
177 - send(json!({"jsonrpc": "2.0", "id": id, "error": {"code": -32601, "message": format!("method not found: {other}")}})).await;
178 - }
179 - }
180 - }
6 + ed_mcp::test_server::run().await;
7 }
src/llm.rs deleted
-507
@@ -1,507 +0,0 @@
1 -//! Minimal streaming client for the Onde Inference (Onde Cloud) OpenAI-compatible
2 -//! `/chat/completions` endpoint.
3 -
4 -use anyhow::{Context, Result, bail};
5 -use futures::StreamExt;
6 -use serde::Deserialize;
7 -use serde_json::{Value, json};
8 -
9 -const DEFAULT_BASE_URL: &str = "https://cloud.ondeinference.com/v1";
10 -const DEFAULT_MODEL: &str = "onde-kkk";
11 -
12 -#[derive(Debug, Clone, PartialEq, Eq)]
13 -pub struct LlmConfig {
14 - pub base_url: String,
15 - /// Onde credentials: `app-id:app-secret`, sent as the bearer token.
16 - pub api_key: Option<String>,
17 - pub model: String,
18 -}
19 -
20 -impl LlmConfig {
21 - /// `ONDE_API_KEY` from the environment, falling back to `config_dir()/env`.
22 - /// `SPLITFIRE_BASE_URL` and `SPLITFIRE_MODEL` override the endpoint and default model.
23 - /// Empty values are ignored: GUI launchers (e.g. Zed via launchd) often export
24 - /// variables set to "", which must not mask the config file.
25 - pub fn from_env() -> Self {
26 - let file_vars = load_config_env_vars();
27 - Self::from_lookup(|name| {
28 - std::env::var(name)
29 - .ok()
30 - .filter(|v| !v.is_empty())
31 - .or_else(|| file_vars.get(name).cloned())
32 - })
33 - }
34 -
35 - /// Build a config from an arbitrary variable lookup (tests, setup).
36 - pub fn from_lookup(lookup: impl Fn(&str) -> Option<String>) -> Self {
37 - let env = |name: &str| lookup(name).filter(|v| !v.is_empty());
38 - Self {
39 - base_url: env("SPLITFIRE_BASE_URL")
40 - .unwrap_or_else(|| DEFAULT_BASE_URL.into())
41 - .trim_end_matches('/')
42 - .to_string(),
43 - api_key: env("ONDE_API_KEY"),
44 - model: env("SPLITFIRE_MODEL").unwrap_or_else(|| DEFAULT_MODEL.into()),
45 - }
46 - }
47 -}
48 -
49 -/// Directory for the stored `env` file:
50 -/// `~/.config/splitfire-agent` on Linux, `~/Library/Application Support/splitfire-agent` on
51 -/// macOS, `%APPDATA%\splitfire-agent` on Windows. `SPLITFIRE_DATA_DIR` overrides it (tests).
52 -/// `None` when the user has no config/home directory.
53 -pub fn config_dir() -> Option<std::path::PathBuf> {
54 - if let Some(dir) = std::env::var_os("SPLITFIRE_DATA_DIR").filter(|v| !v.is_empty()) {
55 - return Some(dir.into());
56 - }
57 - directories::BaseDirs::new().map(|dirs| {
58 - #[cfg(target_os = "linux")]
59 - let base = dirs.config_dir();
60 - #[cfg(not(target_os = "linux"))]
61 - let base = dirs.data_dir();
62 - base.join("splitfire-agent")
63 - })
64 -}
65 -
66 -/// The stored `env` file `--setup` writes and `logout` removes.
67 -pub fn config_file() -> Option<std::path::PathBuf> {
68 - config_dir().map(|d| d.join("env"))
69 -}
70 -
71 -/// Where durable state (sessions, audio scratch) lives. Same directory as the `env` file.
72 -pub fn data_dir() -> std::path::PathBuf {
73 - config_dir().unwrap_or_else(|| std::env::temp_dir().join("splitfire-agent"))
74 -}
75 -
76 -fn load_config_env_vars() -> std::collections::HashMap<String, String> {
77 - let mut vars = std::collections::HashMap::new();
78 - let content = config_file().and_then(|p| std::fs::read_to_string(p).ok());
79 - let Some(content) = content else {
80 - return vars;
81 - };
82 - for line in content.lines() {
83 - let line = line.trim();
84 - if line.is_empty() || line.starts_with('#') {
85 - continue;
86 - }
87 - let line = line.strip_prefix("export ").unwrap_or(line).trim();
88 - if let Some((k, v)) = line.split_once('=') {
89 - let mut v = v.trim();
90 - if v.len() >= 2
91 - && ((v.starts_with('"') && v.ends_with('"'))
92 - || (v.starts_with('\'') && v.ends_with('\'')))
93 - {
94 - v = &v[1..v.len() - 1];
95 - }
96 - vars.insert(k.trim().to_string(), v.to_string());
97 - }
98 - }
99 - vars
100 -}
101 -
102 -#[cfg(test)]
103 -mod tests {
104 - use super::*;
105 -
106 - fn config(vars: &[(&str, &str)]) -> LlmConfig {
107 - LlmConfig::from_lookup(|name| {
108 - vars.iter()
109 - .find(|(k, _)| *k == name)
110 - .map(|(_, v)| v.to_string())
111 - })
112 - }
113 -
114 - #[test]
115 - fn defaults_point_at_onde_cloud() {
116 - let c = config(&[("ONDE_API_KEY", "app:secret")]);
117 - assert_eq!(c.base_url, "https://cloud.ondeinference.com/v1");
118 - assert_eq!(c.api_key.as_deref(), Some("app:secret"));
119 - assert_eq!(c.model, "onde-kkk");
120 - assert_eq!(config(&[]).api_key, None);
121 - assert_eq!(config(&[("ONDE_API_KEY", "")]).api_key, None);
122 - }
123 -
124 - #[test]
125 - fn base_url_and_model_overrides_apply() {
126 - let c = config(&[
127 - ("ONDE_API_KEY", "app:secret"),
128 - ("SPLITFIRE_MODEL", "onde-prism"),
129 - ("SPLITFIRE_BASE_URL", "http://x/v1/"),
130 - ]);
131 - assert_eq!(
132 - (c.base_url.as_str(), c.model.as_str()),
133 - ("http://x/v1", "onde-prism")
134 - );
135 - }
136 -
137 - /// One test touches the process environment, so it can't race another.
138 - #[test]
139 - fn key_from_env_then_file_ignoring_empty_env() {
140 - let dir = std::env::temp_dir().join(format!("sf-llm-{}", uuid::Uuid::new_v4()));
141 - std::fs::create_dir_all(&dir).unwrap();
142 - std::fs::write(
143 - dir.join("env"),
144 - "# comment\nexport ONDE_API_KEY=\"file:key\"\nSPLITFIRE_MODEL=m1\n",
145 - )
146 - .unwrap();
147 - // SAFETY: only this test mutates these variables.
148 - unsafe {
149 - std::env::set_var("SPLITFIRE_DATA_DIR", &dir);
150 - std::env::set_var("ONDE_API_KEY", "");
151 - }
152 - let c = LlmConfig::from_env();
153 - assert_eq!(c.api_key.as_deref(), Some("file:key"));
154 - assert_eq!(c.model, "m1");
155 - unsafe { std::env::set_var("ONDE_API_KEY", "env:key") };
156 - assert_eq!(LlmConfig::from_env().api_key.as_deref(), Some("env:key"));
157 - unsafe {
158 - std::env::remove_var("ONDE_API_KEY");
159 - std::env::remove_var("SPLITFIRE_DATA_DIR");
160 - }
161 - let _ = std::fs::remove_dir_all(dir);
162 - }
163 -
164 - #[test]
165 - fn usage_tolerates_missing_total() {
166 - let chunk: Chunk = serde_json::from_str(
167 - r#"{"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3}}"#,
168 - )
169 - .unwrap();
170 - assert_eq!(chunk.usage.unwrap().used(), 10);
171 - let full = Usage {
172 - prompt_tokens: 1,
173 - completion_tokens: 1,
174 - total_tokens: 5,
175 - };
176 - assert_eq!(full.used(), 5);
177 - }
178 -
179 - #[test]
180 - fn parses_models_response() {
181 - let body = r#"{
182 - "object": "list",
183 - "data": [
184 - {"id": "onde-balanced", "object": "model", "created": 1700000000, "owned_by": "onde"},
185 - {"id": "custom-model"}
186 - ]
187 - }"#;
188 - let parsed: ModelsResponse = serde_json::from_str(body).unwrap();
189 - assert_eq!(parsed.data.len(), 2);
190 - assert_eq!(parsed.data[0].owned_by.as_deref(), Some("onde"));
191 - assert_eq!(parsed.data[1].owned_by, None);
192 - }
193 -}
194 -
195 -#[derive(Debug, Clone, Default)]
196 -pub struct ToolCallRequest {
197 - pub id: String,
198 - pub name: String,
199 - pub arguments: String,
200 -}
201 -
202 -#[derive(Debug, Default)]
203 -pub struct Completion {
204 - pub content: String,
205 - pub tool_calls: Vec<ToolCallRequest>,
206 - pub finish_reason: Option<String>,
207 - /// Token usage, when the endpoint reports it (`stream_options.include_usage`).
208 - pub usage: Option<Usage>,
209 -}
210 -
211 -#[derive(Debug, Clone, Copy, Default, Deserialize)]
212 -pub struct Usage {
213 - #[serde(default)]
214 - pub prompt_tokens: u64,
215 - #[serde(default)]
216 - pub completion_tokens: u64,
217 - #[serde(default)]
218 - pub total_tokens: u64,
219 -}
220 -
221 -impl Usage {
222 - /// Tokens in the context after this completion. Some servers leave `total_tokens` at 0.
223 - pub fn used(&self) -> u64 {
224 - self.total_tokens
225 - .max(self.prompt_tokens + self.completion_tokens)
226 - }
227 -}
228 -
229 -impl Completion {
230 - /// The assistant message to append to the conversation history.
231 - pub fn to_message(&self) -> Value {
232 - let mut msg = json!({ "role": "assistant", "content": self.content });
233 - if !self.tool_calls.is_empty() {
234 - msg["tool_calls"] = self
235 - .tool_calls
236 - .iter()
237 - .map(|tc| {
238 - json!({
239 - "id": tc.id,
240 - "type": "function",
241 - "function": { "name": tc.name, "arguments": tc.arguments },
242 - })
243 - })
244 - .collect();
245 - }
246 - msg
247 - }
248 -}
249 -
250 -/// Streamed pieces surfaced to the caller as they arrive.
251 -pub enum Delta<'a> {
252 - Text(&'a str),
253 - Reasoning(&'a str),
254 -}
255 -
256 -#[derive(Deserialize)]
257 -struct Chunk {
258 - #[serde(default)]
259 - choices: Vec<Choice>,
260 - usage: Option<Usage>,
261 -}
262 -
263 -#[derive(Deserialize)]
264 -struct Choice {
265 - #[serde(default)]
266 - delta: ChunkDelta,
267 - finish_reason: Option<String>,
268 -}
269 -
270 -#[derive(Deserialize, Default)]
271 -struct ChunkDelta {
272 - content: Option<String>,
273 - // Non-standard but common (DeepSeek, vLLM, OpenRouter, llama.cpp).
274 - reasoning_content: Option<String>,
275 - reasoning: Option<String>,
276 - #[serde(default)]
277 - tool_calls: Vec<ToolCallDelta>,
278 -}
279 -
280 -#[derive(Deserialize)]
281 -struct ToolCallDelta {
282 - #[serde(default)]
283 - index: usize,
284 - id: Option<String>,
285 - function: Option<FunctionDelta>,
286 -}
287 -
288 -#[derive(Deserialize)]
289 -struct FunctionDelta {
290 - name: Option<String>,
291 - arguments: Option<String>,
292 -}
293 -
294 -/// A model from `GET /v1/models`. Extra fields in the response are ignored.
295 -#[derive(Debug, Clone, Deserialize)]
296 -pub struct ModelInfo {
297 - pub id: String,
298 - #[serde(default)]
299 - pub owned_by: Option<String>,
300 -}
301 -
302 -#[derive(Deserialize)]
303 -struct ModelsResponse {
304 - data: Vec<ModelInfo>,
305 -}
306 -
307 -#[derive(Debug, Clone)]
308 -pub struct LlmClient {
309 - http: reqwest::Client,
310 - config: LlmConfig,
311 -}
312 -
313 -impl LlmClient {
314 - pub fn new(config: LlmConfig) -> Self {
315 - Self {
316 - http: reqwest::Client::new(),
317 - config,
318 - }
319 - }
320 -
321 - pub fn model(&self) -> &str {
322 - &self.config.model
323 - }
324 -
325 - pub fn config(&self) -> &LlmConfig {
326 - &self.config
327 - }
328 -
329 - /// Whether an API key is configured.
330 - pub fn has_api_key(&self) -> bool {
331 - self.config.api_key.is_some()
332 - }
333 -
334 - /// Verify the configured key with a minimal chat completion.
335 - /// Returns `Err` when no key is set, the key is rejected, or the endpoint is unreachable.
336 - pub async fn check_auth(&self) -> Result<()> {
337 - let Some(key) = &self.config.api_key else {
338 - bail!("no ONDE_API_KEY configured");
339 - };
340 - let body = json!({
341 - "model": self.config.model,
342 - "messages": [{ "role": "user", "content": "ping" }],
343 - "max_tokens": 1,
344 - "stream": false,
345 - });
346 - let resp = self
347 - .http
348 - .post(format!("{}/chat/completions", self.config.base_url))
349 - .json(&body)
350 - .bearer_auth(key)
351 - .send()
352 - .await
353 - .context("reaching Onde Cloud")?;
354 - let status = resp.status();
355 - if matches!(status.as_u16(), 401 | 403) {
356 - bail!("Onde Cloud rejected the API key ({status})");
357 - }
358 - if !status.is_success() {
359 - let text = resp.text().await.unwrap_or_default();
360 - bail!("Onde Cloud returned {status}: {text}");
361 - }
362 - Ok(())
363 - }
364 -
365 - /// List the models the configured endpoint serves (`GET {base_url}/models`).
366 - pub async fn models(&self) -> Result<Vec<ModelInfo>> {
367 - let url = format!("{}/models", self.config.base_url);
368 - let mut req = self.http.get(&url);
369 - if let Some(key) = &self.config.api_key {
370 - req = req.bearer_auth(key);
371 - }
372 - let mut resp = req.send().await.context("sending models request")?;
373 - // Some endpoints (Onde Cloud) list models publicly but reject an invalid key, so a
374 - // bad key shouldn't hide the catalog: retry without auth.
375 - if self.config.api_key.is_some() && matches!(resp.status().as_u16(), 401 | 403) {
376 - resp = self
377 - .http
378 - .get(&url)
379 - .send()
380 - .await
381 - .context("sending models request")?;
382 - }
383 - if !resp.status().is_success() {
384 - let status = resp.status();
385 - let text = resp.text().await.unwrap_or_default();
386 - bail!("models endpoint returned {status}: {text}");
387 - }
388 - let mut models = resp
389 - .json::<ModelsResponse>()
390 - .await
391 - .context("parsing models response")?
392 - .data;
393 - models.sort_by(|a, b| a.id.cmp(&b.id));
394 - Ok(models)
395 - }
396 -
397 - /// Run one streaming chat completion, invoking `on_delta` for every text fragment.
398 - pub async fn complete(
399 - &self,
400 - model: &str,
401 - messages: &[Value],
402 - tools: &Value,
403 - mut on_delta: impl FnMut(Delta<'_>),
404 - ) -> Result<Completion> {
405 - if self.config.api_key.is_none() {
406 - bail!(
407 - "No API key configured. Set ONDE_API_KEY in your environment or run `splitfire-agent --setup`"
408 - );
409 - }
410 - let mut body = json!({
411 - "model": model,
412 - "messages": messages,
413 - "tools": tools,
414 - "stream": true,
415 - "stream_options": { "include_usage": true },
416 - });
417 - let send = |body: &Value| {
418 - let mut req = self
419 - .http
420 - .post(format!("{}/chat/completions", self.config.base_url))
421 - .json(body);
422 - if let Some(key) = &self.config.api_key {
423 - req = req.bearer_auth(key);
424 - }
425 - req.send()
426 - };
427 - let mut resp = send(&body)
428 - .await
429 - .context("sending chat completion request")?;
430 - // Not every OpenAI-compatible server knows `stream_options`; retry without it.
431 - if matches!(resp.status().as_u16(), 400 | 422) {
432 - body.as_object_mut().unwrap().remove("stream_options");
433 - resp = send(&body)
434 - .await
435 - .context("sending chat completion request")?;
436 - }
437 - if !resp.status().is_success() {
438 - let status = resp.status();
439 - let text = resp.text().await.unwrap_or_default();
440 - bail!("LLM endpoint returned {status}: {text}");
441 - }
442 -
443 - let mut out = Completion::default();
444 - let mut stream = resp.bytes_stream();
445 - let mut buf: Vec<u8> = Vec::new();
446 - while let Some(bytes) = stream.next().await {
447 - buf.extend_from_slice(&bytes.context("reading response stream")?);
448 - while let Some(pos) = buf.iter().position(|&b| b == b'\n') {
449 - let line: Vec<u8> = buf.drain(..=pos).collect();
450 - let line = String::from_utf8_lossy(&line);
451 - let Some(data) = line.trim().strip_prefix("data:") else {
452 - continue;
453 - };
454 - let data = data.trim();
455 - if data == "[DONE]" {
456 - return Ok(out);
457 - }
458 - let chunk: Chunk = match serde_json::from_str(data) {
459 - Ok(c) => c,
460 - Err(e) => {
461 - tracing::warn!("skipping unparseable chunk ({e}): {data}");
462 - continue;
463 - }
464 - };
465 - if chunk.usage.is_some() {
466 - out.usage = chunk.usage;
467 - }
468 - for choice in chunk.choices {
469 - let d = choice.delta;
470 - if let Some(r) = d.reasoning_content.as_deref().or(d.reasoning.as_deref())
471 - && !r.is_empty()
472 - {
473 - on_delta(Delta::Reasoning(r));
474 - }
475 - if let Some(t) = d.content.as_deref()
476 - && !t.is_empty()
477 - {
478 - out.content.push_str(t);
479 - on_delta(Delta::Text(t));
480 - }
481 - for tc in d.tool_calls {
482 - if out.tool_calls.len() <= tc.index {
483 - out.tool_calls
484 - .resize(tc.index + 1, ToolCallRequest::default());
485 - }
486 - let slot = &mut out.tool_calls[tc.index];
487 - if let Some(id) = tc.id {
488 - slot.id = id;
489 - }
490 - if let Some(f) = tc.function {
491 - if let Some(n) = f.name {
492 - slot.name.push_str(&n);
493 - }
494 - if let Some(a) = f.arguments {
495 - slot.arguments.push_str(&a);
496 - }
497 - }
498 - }
499 - if choice.finish_reason.is_some() {
500 - out.finish_reason = choice.finish_reason;
501 - }
502 - }
503 - }
504 - }
505 - Ok(out)
506 - }
507 -}
src/main.rs
+24 -402
@@ -1,37 +1,26 @@
1 //! splitfire-agent: a music-domain ACP agent (theory, production, industry, culture, live shows)
2 -//! backed by Onde Inference.
2 +//! backed by Onde Inference, built on the Onde Agent Platform crates: `ed-acp` (the ACP server),
3 +//! `ed-acp-tui` (the terminal UI) and `ed-mcp`. This crate holds only the music: its tools,
4 +//! prompt, slash commands and the [`profile::SplitFire`] that wires them in.
5 //!
6 //! Run from a terminal it opens a TUI; launched by an editor (stdin not a TTY) or with `--acp`
7 //! it speaks ACP over stdio. Set ONDE_API_KEY (`app-id:app-secret`) or run `--setup`.
8 //! SPLITFIRE_YOLO=1 skips permission prompts. Logs go to stderr (RUST_LOG).
9
8 -mod agent;
10 mod audio;
10 -mod llm;
11 -mod mcp;
11 +mod profile;
12 mod prompt;
13 -mod session_store;
13 mod theory;
14 mod tools;
16 -mod tui;
15
16 use std::io::IsTerminal;
17 use std::path::PathBuf;
18 +use std::sync::Arc;
19
21 -use agent_client_protocol::schema::v1::{
22 - AgentAuthCapabilities, AgentCapabilities, AuthenticateRequest, CancelNotification,
23 - CloseSessionRequest, DeleteSessionRequest, Implementation, InitializeRequest,
24 - InitializeResponse, ListSessionsRequest, LoadSessionRequest, LogoutCapabilities, LogoutRequest,
25 - NewSessionRequest, PromptCapabilities, PromptRequest, ResumeSessionRequest,
26 - SessionAdditionalDirectoriesCapabilities, SessionCapabilities, SessionCloseCapabilities,
27 - SessionDeleteCapabilities, SessionListCapabilities, SessionResumeCapabilities,
28 - SetSessionConfigOptionRequest,
29 -};
30 -use agent_client_protocol::{Agent, Stdio};
31 -use anyhow::Context;
20 +use ed_acp::{LlmConfig, ServeOptions, cli};
21 +use ed_acp_tui::TuiConfig;
22
33 -use agent::SplitFireAgent;
34 -use llm::{LlmClient, LlmConfig};
23 +use profile::{INFO, SplitFire};
24
25 const USAGE: &str = "\
26 Usage: splitfire-agent [--acp] [--yolo] [--list-models] [--root <path>...] [--setup] [--version]
@@ -58,7 +47,7 @@ async fn main() -> anyhow::Result<()> {
47 return Ok(());
48 }
49 if has("-V") || has("--version") {
61 - println!("splitfire-agent {}", env!("CARGO_PKG_VERSION"));
50 + println!("{} {}", INFO.name, INFO.version);
51 return Ok(());
52 }
53 // Collect `--root <path>` pairs; validate everything else is a known flag.
@@ -87,398 +76,31 @@ async fn main() -> anyhow::Result<()> {
76 }
77 i += 1;
78 }
90 - let yolo = has("--yolo") || env_flag("SPLITFIRE_YOLO");
79 + let mut opts = ServeOptions::from_env(&INFO);
80 + opts.yolo |= has("--yolo");
81 if has("--setup") {
92 - setup().await?;
82 + cli::setup(&INFO).await?;
83 } else if has("--list-models") {
94 - list_models().await?;
84 + cli::list_models(&INFO).await?;
85 } else if has("--acp") || !std::io::stdin().is_terminal() {
86 if !roots.is_empty() {
87 anyhow::bail!("--root is only supported in the interactive TUI\n\n{USAGE}");
88 }
99 - run_agent(yolo).await?;
89 + cli::init_logging();
90 + ed_acp::serve_stdio(Arc::new(SplitFire), opts).await?;
91 } else {
101 - if !LlmConfig::from_env().api_key.is_some() {
92 + let llm = LlmConfig::from_env(&INFO.env);
93 + if llm.api_key.is_none() {
94 anyhow::bail!(
103 - "No Onde API key configured. Set ONDE_API_KEY or run `splitfire-agent --setup`."
95 + "No Onde API key configured. Set ONDE_API_KEY or run `{} --setup`.",
96 + INFO.name
97 );
98 }
106 - tui::run(yolo, roots).await?;
99 + let mut tui = TuiConfig::current_exe(INFO.name, opts.yolo)?;
100 + tui.env.push((INFO.env.var("SURFACE"), "tui".to_string()));
101 + tui.model_label = llm.model;
102 + tui.extra_roots = roots;
103 + ed_acp_tui::run(tui).await?;
104 }
105 Ok(())
106 }
110 -
111 -fn env_flag(name: &str) -> bool {
112 - std::env::var(name).is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
113 -}
114 -
115 -/// Onde credentials are `app-id:app-secret`: exactly one colon, both halves non-empty.
116 -fn valid_onde_key(key: &str) -> bool {
117 - let mut parts = key.split(':');
118 - matches!(
119 - (parts.next(), parts.next(), parts.next()),
120 - (Some(id), Some(secret), None) if !id.is_empty() && !secret.is_empty()
121 - )
122 -}
123 -
124 -/// `env` file content with `ONDE_API_KEY` set, keeping any other lines.
125 -fn with_api_key(existing: &str, key: &str) -> String {
126 - let mut lines: Vec<String> = existing
127 - .lines()
128 - .filter(|l| {
129 - let l = l.trim().strip_prefix("export ").unwrap_or(l.trim());
130 - !l.starts_with("ONDE_API_KEY=")
131 - })
132 - .map(String::from)
133 - .collect();
134 - lines.push(format!("ONDE_API_KEY={key}"));
135 - lines.join("\n") + "\n"
136 -}
137 -
138 -/// Interactive first-run setup (`--setup`): prompt for the Onde key, verify it against Onde
139 -/// Cloud, and store it in the platform config dir (`config_dir()/env`, mode 0600).
140 -async fn setup() -> anyhow::Result<()> {
141 - use std::io::Write;
142 -
143 - if !std::io::stdin().is_terminal() {
144 - anyhow::bail!("--setup needs an interactive terminal");
145 - }
146 - println!("splitfire-agent setup\n");
147 - println!("Get credentials: sign in at https://ondeinference.com/root/login,");
148 - println!("register an app and assign a model. Your key is \"app-id:app-secret\".\n");
149 - let key = loop {
150 - print!("Paste your ONDE_API_KEY: ");
151 - std::io::stdout().flush()?;
152 - let mut line = String::new();
153 - if std::io::stdin().read_line(&mut line)? == 0 {
154 - anyhow::bail!("no input");
155 - }
156 - let key = line.trim().to_string();
157 - if valid_onde_key(&key) {
158 - break key;
159 - }
160 - println!("Onde credentials look like \"app-id:app-secret\" (exactly one colon).");
161 - };
162 -
163 - // Verify before writing anything.
164 - let mut config = LlmConfig::from_env();
165 - config.api_key = Some(key.clone());
166 - print!("Verifying the key with Onde Cloud… ");
167 - std::io::stdout().flush()?;
168 - match LlmClient::new(config).check_auth().await {
169 - Ok(()) => println!("OK"),
170 - Err(e) => anyhow::bail!("\nKey check failed: {e:#}\nNothing was written."),
171 - }
172 -
173 - let dir = llm::config_dir().context("no config directory available")?;
174 - std::fs::create_dir_all(&dir)?;
175 - let path = dir.join("env");
176 - let existing = std::fs::read_to_string(&path).unwrap_or_default();
177 - std::fs::write(&path, with_api_key(&existing, &key))?;
178 - #[cfg(unix)]
179 - {
180 - use std::os::unix::fs::PermissionsExt;
181 - std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;
182 - }
183 - println!("\nWrote {}. You're all set.", path.display());
184 - Ok(())
185 -}
186 -
187 -async fn list_models() -> anyhow::Result<()> {
188 - let client = LlmClient::new(LlmConfig::from_env());
189 - for m in client.models().await? {
190 - match m.owned_by {
191 - Some(owner) => println!("{}\t{}", m.id, owner),
192 - None => println!("{}", m.id),
193 - }
194 - }
195 - Ok(())
196 -}
197 -
198 -/// How the agent was launched: `tui` (interactive terminal UI) or `acp` (editor).
199 -/// The TUI sets `SPLITFIRE_SURFACE=tui` on its subprocess; editors launch `--acp`
200 -/// directly so the default is `acp`.
201 -fn surface() -> &'static str {
202 - match std::env::var("SPLITFIRE_SURFACE") {
203 - Ok(s) if s == "tui" => "tui",
204 - _ => "acp",
205 - }
206 -}
207 -
208 -async fn run_agent(yolo: bool) -> agent_client_protocol::Result<()> {
209 - tracing_subscriber::fmt()
210 - .with_writer(std::io::stderr)
211 - .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
212 - .init();
213 -
214 - let agent = SplitFireAgent::new(LlmConfig::from_env(), yolo, surface(), llm::data_dir());
215 - tracing::info!("splitfire-agent starting with model {}", agent.model());
216 -
217 - let result = Agent
218 - .builder()
219 - .name("splitfire-agent")
220 - .on_receive_request(
221 - {
222 - let agent = agent.clone();
223 - async move |req: InitializeRequest, responder, _cx| {
224 - agent.set_client_caps(req.client_capabilities.clone());
225 - // Terminal auth is only usable by clients that can run it.
226 - let methods = if req.client_capabilities.auth.terminal {
227 - agent::auth_methods()
228 - } else {
229 - Vec::new()
230 - };
231 - responder.respond(
232 - InitializeResponse::new(agent::negotiate_version(req.protocol_version))
233 - .auth_methods(methods)
234 - .agent_capabilities(
235 - AgentCapabilities::new()
236 - .prompt_capabilities(
237 - PromptCapabilities::new()
238 - .embedded_context(true)
239 - .image(true)
240 - .audio(true),
241 - )
242 - .session_capabilities(
243 - SessionCapabilities::new()
244 - .list(SessionListCapabilities::new())
245 - .additional_directories(
246 - SessionAdditionalDirectoriesCapabilities::new(),
247 - )
248 - .resume(SessionResumeCapabilities::new())
249 - .close(SessionCloseCapabilities::new())
250 - .delete(SessionDeleteCapabilities::new()),
251 - )
252 - .auth(
253 - AgentAuthCapabilities::new()
254 - .logout(LogoutCapabilities::new()),
255 - )
256 - .load_session(true),
257 - )
258 - .agent_info(Implementation::new(
259 - "splitfire-agent",
260 - env!("CARGO_PKG_VERSION"),
261 - )),
262 - )
263 - }
264 - },
265 - agent_client_protocol::on_receive_request!(),
266 - )
267 - .on_receive_request(
268 - {
269 - let agent = agent.clone();
270 - async move |req: AuthenticateRequest, responder, cx| {
271 - let agent = agent.clone();
272 - cx.spawn(async move {
273 - match agent.authenticate(req).await {
274 - Ok(resp) => responder.respond(resp),
275 - Err(e) => responder.respond_with_error(e),
276 - }
277 - })
278 - }
279 - },
280 - agent_client_protocol::on_receive_request!(),
281 - )
282 - .on_receive_request(
283 - {
284 - let agent = agent.clone();
285 - async move |_req: LogoutRequest, responder, cx| {
286 - let agent = agent.clone();
287 - cx.spawn(async move {
288 - match agent.logout().await {
289 - Ok(resp) => responder.respond(resp),
290 - Err(e) => responder.respond_with_error(e),
291 - }
292 - })
293 - }
294 - },
295 - agent_client_protocol::on_receive_request!(),
296 - )
297 - .on_receive_request(
298 - {
299 - let agent = agent.clone();
300 - async move |req: NewSessionRequest, responder, cx| {
301 - // Listing models and starting MCP servers are slow; keep them off the
302 - // dispatch loop.
303 - let agent = agent.clone();
304 - let connection = cx.clone();
305 - cx.spawn(async move {
306 - // Picks up a key stored by `--setup` since this process started.
307 - if !agent.refresh_credentials().await {
308 - return responder.respond_with_error(agent::auth_required_error());
309 - }
310 - match agent.new_session(req).await {
311 - Ok(resp) => {
312 - let id = resp.session_id.clone();
313 - responder.respond(resp)?;
314 - agent.send_commands(&connection, &id)
315 - }
316 - Err(e) => responder.respond_with_error(e),
317 - }
318 - })
319 - }
320 - },
321 - agent_client_protocol::on_receive_request!(),
322 - )
323 - .on_receive_request(
324 - {
325 - let agent = agent.clone();
326 - async move |req: ListSessionsRequest, responder, cx| {
327 - // Reads the session directory from disk.
328 - let agent = agent.clone();
329 - cx.spawn(async move {
330 - match tokio::task::spawn_blocking(move || agent.list_sessions(req)).await {
331 - Ok(Ok(resp)) => responder.respond(resp),
332 - Ok(Err(e)) => responder.respond_with_error(e),
333 - Err(e) => responder.respond_with_internal_error(e.to_string()),
334 - }
335 - })
336 - }
337 - },
338 - agent_client_protocol::on_receive_request!(),
339 - )
340 - .on_receive_request(
341 - {
342 - let agent = agent.clone();
343 - async move |req: LoadSessionRequest, responder, cx| {
344 - let agent = agent.clone();
345 - let connection = cx.clone();
346 - cx.spawn(async move {
347 - // Editors reopen a thread with session/load after terminal auth, so
348 - // AUTH_REQUIRED here shows sign-in instead of a dead thread.
349 - if !agent.refresh_credentials().await {
350 - return responder.respond_with_error(agent::auth_required_error());
351 - }
352 - let id = req.session_id.clone();
353 - match agent.load_session(req, &connection).await {
354 - Ok(resp) => {
355 - responder.respond(resp)?;
356 - agent.send_commands(&connection, &id)
357 - }
358 - Err(e) => responder.respond_with_error(e),
359 - }
360 - })
361 - }
362 - },
363 - agent_client_protocol::on_receive_request!(),
364 - )
365 - .on_receive_request(
366 - {
367 - let agent = agent.clone();
368 - async move |req: ResumeSessionRequest, responder, cx| {
369 - let agent = agent.clone();
370 - let connection = cx.clone();
371 - cx.spawn(async move {
372 - if !agent.refresh_credentials().await {
373 - return responder.respond_with_error(agent::auth_required_error());
374 - }
375 - let id = req.session_id.clone();
376 - match agent.resume_session(req, &connection).await {
377 - Ok(resp) => {
378 - responder.respond(resp)?;
379 - agent.send_commands(&connection, &id)
380 - }
381 - Err(e) => responder.respond_with_error(e),
382 - }
383 - })
384 - }
385 - },
386 - agent_client_protocol::on_receive_request!(),
387 - )
388 - .on_receive_request(
389 - {
390 - let agent = agent.clone();
391 - async move |req: CloseSessionRequest, responder, _cx| match agent.close_session(req)
392 - {
393 - Ok(resp) => responder.respond(resp),
394 - Err(e) => responder.respond_with_error(e),
395 - }
396 - },
397 - agent_client_protocol::on_receive_request!(),
398 - )
399 - .on_receive_request(
400 - {
401 - let agent = agent.clone();
402 - async move |req: DeleteSessionRequest, responder, _cx| match agent
403 - .delete_session(req)
404 - {
405 - Ok(resp) => responder.respond(resp),
406 - Err(e) => responder.respond_with_error(e),
407 - }
408 - },
409 - agent_client_protocol::on_receive_request!(),
410 - )
411 - .on_receive_request(
412 - {
413 - let agent = agent.clone();
414 - async move |req: SetSessionConfigOptionRequest, responder, cx| {
415 - let agent = agent.clone();
416 - let connection = cx.clone();
417 - cx.spawn(async move {
418 - match agent.set_config_option(req, &connection).await {
419 - Ok(resp) => responder.respond(resp),
420 - Err(e) => responder.respond_with_error(e),
421 - }
422 - })
423 - }
424 - },
425 - agent_client_protocol::on_receive_request!(),
426 - )
427 - .on_receive_request(
428 - {
429 - let agent = agent.clone();
430 - async move |req: PromptRequest, responder, cx| {
431 - // Run the turn off the dispatch loop so it can make requests to the client.
432 - let agent = agent.clone();
433 - let connection = cx.clone();
434 - cx.spawn(async move {
435 - if !agent.has_api_key() && !agent.refresh_credentials().await {
436 - return responder.respond_with_error(agent::auth_required_error());
437 - }
438 - agent.prompt(req, responder, connection).await
439 - })
440 - }
441 - },
442 - agent_client_protocol::on_receive_request!(),
443 - )
444 - .on_receive_notification(
445 - {
446 - let agent = agent.clone();
447 - async move |n: CancelNotification, _cx| {
448 - agent.cancel(&n.session_id);
449 - Ok(())
450 - }
451 - },
452 - agent_client_protocol::on_receive_notification!(),
453 - )
454 - .connect_to(Stdio::new())
455 - .await;
456 - // The editor went away: stop MCP servers instead of leaving them to be orphaned.
457 - agent.shutdown().await;
458 - result
459 -}
460 -
461 -#[cfg(test)]
462 -mod tests {
463 - use super::*;
464 -
465 - #[test]
466 - fn onde_key_shape() {
467 - assert!(valid_onde_key("app:secret"));
468 - assert!(!valid_onde_key("secret"));
469 - assert!(!valid_onde_key("a:b:c"));
470 - assert!(!valid_onde_key(":secret"));
471 - assert!(!valid_onde_key("app:"));
472 - assert!(!valid_onde_key(""));
473 - }
474 -
475 - #[test]
476 - fn env_file_keeps_other_lines_and_replaces_key() {
477 - let out = with_api_key(
478 - "# hi\nSPLITFIRE_MODEL=m\nexport ONDE_API_KEY=old:old\n",
479 - "new:key",
480 - );
481 - assert_eq!(out, "# hi\nSPLITFIRE_MODEL=m\nONDE_API_KEY=new:key\n");
482 - assert_eq!(with_api_key("", "a:b"), "ONDE_API_KEY=a:b\n");
483 - }
484 -}
src/mcp.rs deleted
-553
@@ -1,553 +0,0 @@
1 -//! MCP client over stdio (rmcp): servers supplied by the ACP client in `session/new|load|resume`,
2 -//! plus the optional built-in `demucs --mcp` stem-separation server.
3 -//!
4 -//! Tools are exposed to the model as `mcp__<server>__<tool>`. A failing server is never fatal:
5 -//! it is logged and its tools are simply absent.
6 -
7 -use std::collections::HashMap;
8 -use std::path::{Path, PathBuf};
9 -use std::sync::Arc;
10 -use std::time::Duration;
11 -
12 -use agent_client_protocol::schema::v1::{McpServer, ToolCallContent, ToolCallUpdateFields};
13 -use anyhow::{Context, Result, anyhow, bail};
14 -use futures::StreamExt;
15 -use rmcp::handler::client::progress::ProgressDispatcher;
16 -use rmcp::model::{
17 - CallToolRequestParams, ClientCapabilities, ClientRequest, ContentBlock, Implementation,
18 - ProgressNotificationParam, Request, ServerResult, Tool,
19 -};
20 -use rmcp::service::{NotificationContext, PeerRequestOptions, RunningService};
21 -use rmcp::transport::TokioChildProcess;
22 -use rmcp::{ClientHandler, RoleClient, ServiceExt};
23 -use serde_json::{Map, Value, json};
24 -use tokio::sync::Mutex;
25 -
26 -use crate::tools::{ToolCtx, ToolOutcome, absolutize, function_def, truncate};
27 -
28 -const PREFIX: &str = "mcp__";
29 -/// Server name of the built-in stem separator; a client-supplied server of this name wins.
30 -pub const DEMUCS: &str = "demucs";
31 -const CONNECT_TIMEOUT: Duration = Duration::from_secs(20);
32 -/// Upper bound for one tool call: separation of a long track can take many minutes.
33 -const CALL_TIMEOUT: Duration = Duration::from_secs(60 * 60);
34 -const MAX_TOOL_NAME: usize = 64;
35 -
36 -/// Split `mcp__<server>__<tool>` into its parts.
37 -pub fn split_qualified(name: &str) -> Option<(&str, &str)> {
38 - let rest = name.strip_prefix(PREFIX)?;
39 - let (server, tool) = rest.split_once("__")?;
40 - (!server.is_empty() && !tool.is_empty()).then_some((server, tool))
41 -}
42 -
43 -fn sanitize(part: &str) -> String {
44 - let mut out: String = part
45 - .chars()
46 - .map(|c| {
47 - if c.is_ascii_alphanumeric() || c == '-' {
48 - c
49 - } else {
50 - '_'
51 - }
52 - })
53 - .collect();
54 - while out.contains("__") {
55 - out = out.replace("__", "_");
56 - }
57 - out.trim_matches('_').to_string()
58 -}
59 -
60 -fn qualified(server: &str, tool: &str) -> String {
61 - let mut name = format!("{PREFIX}{}__{}", sanitize(server), sanitize(tool));
62 - name.truncate(MAX_TOOL_NAME);
63 - name
64 -}
65 -
66 -struct McpClient {
67 - progress: ProgressDispatcher,
68 -}
69 -
70 -impl ClientHandler for McpClient {
71 - async fn on_progress(
72 - &self,
73 - params: ProgressNotificationParam,
74 - _context: NotificationContext<RoleClient>,
75 - ) {
76 - self.progress.handle_notification(params).await;
77 - }
78 -
79 - fn get_info(&self) -> rmcp::model::ClientConfig {
80 - rmcp::model::ClientConfig::new(
81 - ClientCapabilities::default(),
82 - Implementation::new("splitfire-agent", env!("CARGO_PKG_VERSION")),
83 - )
84 - }
85 -}
86 -
87 -/// One live connection to an MCP server process.
88 -pub struct McpConn {
89 - name: String,
90 - service: RunningService<RoleClient, McpClient>,
91 - tools: Vec<Tool>,
92 -}
93 -
94 -impl McpConn {
95 - async fn connect(
96 - name: &str,
97 - command: &Path,
98 - args: &[String],
99 - env: &[(String, String)],
100 - ) -> Result<McpConn> {
101 - let mut cmd = tokio::process::Command::new(command);
102 - cmd.args(args).envs(env.iter().map(|(k, v)| (k, v)));
103 - let transport = TokioChildProcess::new(cmd)
104 - .with_context(|| format!("starting {}", command.display()))?;
105 - let client = McpClient {
106 - progress: ProgressDispatcher::new(),
107 - };
108 - let service = tokio::time::timeout(CONNECT_TIMEOUT, client.serve(transport))
109 - .await
110 - .map_err(|_| anyhow!("no MCP handshake within {}s", CONNECT_TIMEOUT.as_secs()))?
111 - .map_err(|e| anyhow!("MCP handshake failed: {e}"))?;
112 - let tools = tokio::time::timeout(CONNECT_TIMEOUT, service.peer().list_all_tools())
113 - .await
114 - .map_err(|_| anyhow!("tools/list timed out"))?
115 - .map_err(|e| anyhow!("tools/list failed: {e}"))?;
116 - Ok(McpConn {
117 - name: name.to_string(),
118 - service,
119 - tools,
120 - })
121 - }
122 -
123 - async fn close(mut self) {
124 - let _ = self
125 - .service
126 - .close_with_timeout(Duration::from_secs(2))
127 - .await;
128 - }
129 -}
130 -
131 -struct Entry {
132 - conn: Arc<McpConn>,
133 - raw: String,
134 - description: String,
135 - schema: Value,
136 - read_only: bool,
137 -}
138 -
139 -/// The MCP tools available to one session.
140 -#[derive(Default)]
141 -pub struct McpToolset {
142 - tools: HashMap<String, Entry>,
143 - conns: Vec<Arc<McpConn>>,
144 -}
145 -
146 -/// The process-wide `demucs --mcp` server: started at most once, on first need.
147 -enum DemucsState {
148 - NotStarted,
149 - Unavailable,
150 - Running(Arc<McpConn>),
151 -}
152 -
153 -static DEMUCS_CONN: Mutex<DemucsState> = Mutex::const_new(DemucsState::NotStarted);
154 -
155 -/// `SPLITFIRE_DEMUCS_BIN`, else `demucs` on `PATH`.
156 -fn demucs_binary() -> Option<PathBuf> {
157 - if let Some(bin) = std::env::var_os("SPLITFIRE_DEMUCS_BIN").filter(|v| !v.is_empty()) {
158 - return Some(bin.into());
159 - }
160 - let path = std::env::var_os("PATH")?;
161 - std::env::split_paths(&path)
162 - .map(|dir| dir.join("demucs"))
163 - .find(|candidate| candidate.is_file())
164 -}
165 -
166 -fn demucs_disabled() -> bool {
167 - std::env::var("SPLITFIRE_MCP_DEMUCS").is_ok_and(|v| {
168 - matches!(
169 - v.to_ascii_lowercase().as_str(),
170 - "off" | "0" | "false" | "no"
171 - )
172 - })
173 -}
174 -
175 -/// Start `demucs --mcp` once per process, on first need. `None` (logged once) when it is
176 -/// disabled, missing, or fails its handshake.
177 -async fn demucs() -> Option<Arc<McpConn>> {
178 - let mut state = DEMUCS_CONN.lock().await;
179 - match &*state {
180 - DemucsState::Running(conn) => return Some(conn.clone()),
181 - DemucsState::Unavailable => return None,
182 - DemucsState::NotStarted => {}
183 - }
184 - *state = DemucsState::Unavailable;
185 - if demucs_disabled() {
186 - tracing::info!("stem separation disabled (SPLITFIRE_MCP_DEMUCS=off)");
187 - return None;
188 - }
189 - let Some(bin) = demucs_binary() else {
190 - tracing::info!(
191 - "stem separation unavailable: `demucs` not found (install it or set SPLITFIRE_DEMUCS_BIN)"
192 - );
193 - return None;
194 - };
195 - match McpConn::connect(DEMUCS, &bin, &["--mcp".to_string()], &[]).await {
196 - Ok(conn) => {
197 - tracing::info!("stem separation available via {}", bin.display());
198 - let conn = Arc::new(conn);
199 - *state = DemucsState::Running(conn.clone());
200 - Some(conn)
201 - }
202 - Err(e) => {
203 - tracing::info!(
204 - "stem separation unavailable: {} --mcp failed: {e:#}",
205 - bin.display()
206 - );
207 - None
208 - }
209 - }
210 -}
211 -
212 -/// Stop the process-wide demucs server, if one was started. Call after every session's
213 -/// toolset has been shut down so this is the last owner.
214 -pub async fn shutdown_demucs() {
215 - let state = std::mem::replace(&mut *DEMUCS_CONN.lock().await, DemucsState::Unavailable);
216 - if let DemucsState::Running(conn) = state
217 - && let Some(conn) = Arc::into_inner(conn)
218 - {
219 - conn.close().await;
220 - }
221 -}
222 -
223 -impl McpToolset {
224 - /// Connect the client-supplied stdio servers (HTTP/SSE are not advertised, so ignored)
225 - /// and, unless the client supplied its own `demucs`, the built-in separator.
226 - pub async fn connect(servers: &[McpServer]) -> Self {
227 - let mut set = Self::default();
228 - let mut supplied_demucs = false;
229 - for server in servers {
230 - let McpServer::Stdio(s) = server else {
231 - tracing::info!("ignoring non-stdio MCP server");
232 - continue;
233 - };
234 - supplied_demucs |= s.name == DEMUCS;
235 - let env: Vec<(String, String)> = s
236 - .env
237 - .iter()
238 - .map(|e| (e.name.clone(), e.value.clone()))
239 - .collect();
240 - match McpConn::connect(&s.name, &s.command, &s.args, &env).await {
241 - Ok(conn) => set.add(Arc::new(conn)),
242 - Err(e) => tracing::warn!("MCP server `{}` unavailable: {e:#}", s.name),
243 - }
244 - }
245 - if !supplied_demucs && let Some(conn) = demucs().await {
246 - set.add(conn);
247 - }
248 - set
249 - }
250 -
251 - fn add(&mut self, conn: Arc<McpConn>) {
252 - for tool in &conn.tools {
253 - let name = qualified(&conn.name, &tool.name);
254 - if self.tools.contains_key(&name) {
255 - tracing::warn!("duplicate MCP tool name {name}; keeping the first");
256 - continue;
257 - }
258 - let mut schema = Value::Object((*tool.input_schema).clone());
259 - if schema.get("type").is_none() {
260 - schema["type"] = json!("object");
261 - }
262 - let read_only = tool
263 - .annotations
264 - .as_ref()
265 - .and_then(|a| a.read_only_hint)
266 - .unwrap_or(false);
267 - self.tools.insert(
268 - name,
269 - Entry {
270 - conn: conn.clone(),
271 - raw: tool.name.to_string(),
272 - description: format!(
273 - "[{} MCP server] {}",
274 - conn.name,
275 - tool.description.as_deref().unwrap_or("")
276 - ),
277 - schema,
278 - read_only,
279 - },
280 - );
281 - }
282 - self.conns.push(conn);
283 - }
284 -
285 - pub fn definitions(&self) -> Vec<Value> {
286 - let mut names: Vec<&String> = self.tools.keys().collect();
287 - names.sort();
288 - names
289 - .into_iter()
290 - .map(|n| {
291 - let e = &self.tools[n];
292 - function_def(n, &e.description, e.schema.clone())
293 - })
294 - .collect()
295 - }
296 -
297 - pub fn has(&self, name: &str) -> bool {
298 - self.tools.contains_key(name)
299 - }
300 -
301 - pub fn is_read_only(&self, name: &str) -> bool {
302 - self.tools.get(name).is_some_and(|e| e.read_only)
303 - }
304 -
305 - /// Whether the built-in (or a client-supplied) stem separator is connected.
306 - pub fn has_stem_separation(&self) -> bool {
307 - self.tools
308 - .contains_key(&qualified(DEMUCS, "separate_stems"))
309 - }
310 -
311 - /// What the permission prompt shows for a call. For stem separation this consults
312 - /// `list_models` so the prompt can state the download size of uncached weights.
313 - pub async fn permission_preview(&self, name: &str, args: &Value) -> Vec<ToolCallContent> {
314 - let Some(entry) = self.tools.get(name) else {
315 - return Vec::new();
316 - };
317 - let mut text = format!("Run `{}` on the {} MCP server.", entry.raw, entry.conn.name);
318 - if entry.raw == "separate_stems" && entry.conn.name == DEMUCS {
319 - let input = args.get("input").and_then(Value::as_str).unwrap_or("?");
320 - let model = args
321 - .get("model")
322 - .and_then(Value::as_str)
323 - .unwrap_or("htdemucs");
324 - text = format!(
325 - "Separate {input} into stems with {model}. Runs locally and can take minutes."
326 - );
327 - match model_download_note(entry, model).await {
328 - Some(note) => text.push_str(&format!("\n{note}")),
329 - None => text.push_str("\nCould not check whether the model weights are cached."),
330 - }
331 - } else if args.as_object().is_some_and(|o| !o.is_empty()) {
332 - let shown = serde_json::to_string_pretty(args).unwrap_or_default();
333 - text.push_str(&format!("\nArguments:\n```json\n{}\n```", truncate(shown)));
334 - }
335 - vec![ToolCallContent::from(text)]
336 - }
337 -
338 - /// Call a tool, streaming MCP progress into the tool-call card and forwarding
339 - /// session cancellation as `notifications/cancelled`.
340 - pub async fn call(
341 - &self,
342 - ctx: &ToolCtx,
343 - tool_call_id: &str,
344 - name: &str,
345 - args: Value,
346 - ) -> Result<ToolOutcome> {
347 - let entry = self
348 - .tools
349 - .get(name)
350 - .ok_or_else(|| anyhow!("unknown tool `{name}`"))?;
351 - let arguments: Map<String, Value> = match args {
352 - Value::Object(o) => o,
353 - Value::Null => Map::new(),
354 - _ => bail!("tool arguments must be a JSON object"),
355 - };
356 - let params = CallToolRequestParams::new(entry.raw.clone()).with_arguments(arguments);
357 - let mut handle = entry
358 - .conn
359 - .service
360 - .peer()
361 - .send_cancellable_request(
362 - ClientRequest::CallToolRequest(Request::new(params)),
363 - PeerRequestOptions::no_options(),
364 - )
365 - .await
366 - .map_err(|e| anyhow!("MCP server `{}` is not responding: {e}", entry.conn.name))?;
367 - let mut progress = entry
368 - .conn
369 - .service
370 - .service()
371 - .progress
372 - .subscribe(handle.progress_token.clone())
373 - .await;
374 - let deadline = tokio::time::sleep(CALL_TIMEOUT);
375 - tokio::pin!(deadline);
376 -
377 - enum End {
378 - Done(Box<Result<ServerResult, rmcp::ServiceError>>),
379 - Cancelled,
380 - TimedOut,
381 - }
382 - let end = loop {
383 - tokio::select! {
384 - r = &mut handle.rx => {
385 - break End::Done(Box::new(r.unwrap_or(Err(rmcp::ServiceError::TransportClosed))));
386 - }
387 - Some(p) = progress.next() => {
388 - let line = progress_line(&p);
389 - let _ = ctx.update(
390 - tool_call_id,
391 - ToolCallUpdateFields::new().content(vec![ToolCallContent::from(line)]),
392 - );
393 - }
394 - () = ctx.cancel.cancelled() => break End::Cancelled,
395 - () = &mut deadline => break End::TimedOut,
396 - }
397 - };
398 - match end {
399 - End::Done(done) => match *done {
400 - Ok(ServerResult::CallToolResult(r)) => Ok(outcome_from(r)),
401 - Ok(_) => Ok(ToolOutcome::err(
402 - "The MCP server sent an unexpected response.",
403 - )),
404 - Err(e) => Ok(ToolOutcome::err(format!("MCP call failed: {e}"))),
405 - },
406 - End::Cancelled => {
407 - let _ = handle.cancel(Some("cancelled by user".into())).await;
408 - Ok(ToolOutcome::err("Cancelled by user."))
409 - }
410 - End::TimedOut => {
411 - let _ = handle.cancel(Some("timed out".into())).await;
412 - Ok(ToolOutcome::err(format!(
413 - "Timed out after {} minutes.",
414 - CALL_TIMEOUT.as_secs() / 60
415 - )))
416 - }
417 - }
418 - }
419 -
420 - /// Close this session's servers. Shared servers (the built-in demucs) stay up.
421 - pub async fn shutdown(self) {
422 - drop(self.tools);
423 - for conn in self.conns {
424 - if let Some(conn) = Arc::into_inner(conn) {
425 - conn.close().await;
426 - }
427 - }
428 - }
429 -}
430 -
431 -fn progress_line(p: &ProgressNotificationParam) -> String {
432 - let pct = p
433 - .total
434 - .filter(|t| *t > 0.0)
435 - .map(|t| format!(" ({:.0}%)", 100.0 * p.progress / t));
436 - format!(
437 - "{}{}",
438 - p.message.as_deref().unwrap_or("Working…"),
439 - pct.unwrap_or_default()
440 - )
441 -}
442 -
443 -fn outcome_from(r: rmcp::model::CallToolResult) -> ToolOutcome {
444 - let mut parts = Vec::new();
445 - for block in &r.content {
446 - match block {
447 - ContentBlock::Text(t) => parts.push(t.text.clone()),
448 - ContentBlock::Image(i) => parts.push(format!("[image: {}]", i.mime_type)),
449 - ContentBlock::Audio(a) => parts.push(format!("[audio: {}]", a.mime_type)),
450 - ContentBlock::ResourceLink(l) => parts.push(format!("[resource: {}]", l.uri)),
451 - ContentBlock::Resource(_) => parts.push("[embedded resource]".into()),
452 - _ => parts.push("[unsupported content]".into()),
453 - }
454 - }
455 - let mut text = parts.join("\n");
456 - if text.is_empty() {
457 - text = r
458 - .structured_content
459 - .as_ref()
460 - .map(|v| v.to_string())
461 - .unwrap_or_else(|| "(no output)".into());
462 - }
463 - let mut out = if r.is_error == Some(true) {
464 - ToolOutcome::err(truncate(text))
465 - } else {
466 - ToolOutcome::ok(truncate(text))
467 - };
468 - out.locations = stem_paths(r.structured_content.as_ref());
469 - out
470 -}
471 -
472 -/// `structuredContent.stems[].path`, absolute paths only.
473 -fn stem_paths(structured: Option<&Value>) -> Vec<PathBuf> {
474 - structured
475 - .and_then(|v| v.get("stems"))
476 - .and_then(Value::as_array)
477 - .map(|stems| {
478 - stems
479 - .iter()
480 - .filter_map(|s| s.get("path").and_then(Value::as_str))
481 - .map(PathBuf::from)
482 - .filter(|p| p.is_absolute())
483 - .map(|p| absolutize(&p))
484 - .collect()
485 - })
486 - .unwrap_or_default()
487 -}
488 -
489 -/// "Downloads N MB…" when `model` is not cached, per demucs's `list_models`.
490 -async fn model_download_note(entry: &Entry, model: &str) -> Option<String> {
491 - let params = CallToolRequestParams::new("list_models");
492 - let result = tokio::time::timeout(
493 - Duration::from_secs(10),
494 - entry.conn.service.peer().call_tool(params),
495 - )
496 - .await
497 - .ok()?
498 - .ok()?;
499 - let models = result
500 - .structured_content?
501 - .get("models")?
502 - .as_array()?
503 - .clone();
504 - let m = models
505 - .iter()
506 - .find(|m| m.get("id").and_then(Value::as_str) == Some(model))?;
507 - let size = m.get("size_mb").and_then(Value::as_u64);
508 - Some(match (m.get("cached").and_then(Value::as_bool), size) {
509 - (Some(true), _) => "The model weights are already downloaded.".to_string(),
510 - (_, Some(mb)) => format!("First use downloads about {mb} MB of model weights."),
511 - _ => "The model weights may need to be downloaded first.".to_string(),
512 - })
513 -}
514 -
515 -#[cfg(test)]
516 -mod tests {
517 - use super::*;
518 -
519 - #[test]
520 - fn qualified_names_round_trip() {
521 - let q = qualified("My Server!", "do.thing");
522 - assert_eq!(q, "mcp__My_Server__do_thing");
523 - assert_eq!(split_qualified(&q), Some(("My_Server", "do_thing")));
524 - assert_eq!(split_qualified("read_file"), None);
525 - assert_eq!(split_qualified("mcp__x"), None);
526 - assert_eq!(split_qualified("mcp____t"), None);
527 - assert!(qualified("s", &"t".repeat(100)).len() <= MAX_TOOL_NAME);
528 - }
529 -
530 - #[test]
531 - fn stems_become_absolute_locations() {
532 - let v = json!({"stems": [
533 - {"id": "drums", "path": "/tmp/a/drums.wav"},
534 - {"id": "bass", "path": "relative/bass.wav"},
535 - {"id": "x"}
536 - ]});
537 - let paths = stem_paths(Some(&v));
538 - assert_eq!(paths.len(), 1);
539 - assert!(paths[0].ends_with("drums.wav"));
540 - assert!(stem_paths(None).is_empty());
541 - }
542 -
543 - #[test]
544 - fn progress_text() {
545 - let p = ProgressNotificationParam::new(
546 - rmcp::model::ProgressToken(rmcp::model::NumberOrString::Number(1)),
547 - 25.0,
548 - )
549 - .with_total(100.0)
550 - .with_message("Separating");
551 - assert_eq!(progress_line(&p), "Separating (25%)");
552 - }
553 -}
src/profile.rs new
+224
@@ -0,0 +1,224 @@
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 +}
src/prompt.rs
+2 -381
@@ -1,17 +1,10 @@
1 -//! The system prompt, slash commands, and the conversion of ACP prompt content into
2 -//! model content.
1 +//! SplitFire's system prompt and slash commands. Converting prompt content is `ed-acp`'s job.
2
3 use std::path::{Path, PathBuf};
4
5 use agent_client_protocol::schema::v1::{
7 - AvailableCommand, AvailableCommandInput, ClientCapabilities, ContentBlock, EmbeddedResource,
8 - EmbeddedResourceResource, ReadTextFileRequest, ResourceLink, SessionId, TextResourceContents,
9 - UnstructuredCommandInput,
6 + AvailableCommand, AvailableCommandInput, UnstructuredCommandInput,
7 };
11 -use agent_client_protocol::{Client, ConnectionTo};
12 -use anyhow::{Context, Result};
13 -use base64::Engine;
14 -use serde_json::{Value, json};
8
9 pub fn system_prompt(
10 surface: &str,
@@ -191,343 +184,10 @@ pub fn expand_slash(text: &str, stems_available: bool) -> Option<String> {
184 })
185 }
186
194 -fn embedded_text(res: &EmbeddedResource) -> Option<String> {
195 - match &res.resource {
196 - EmbeddedResourceResource::TextResourceContents(r) => {
197 - Some(format!("<file uri=\"{}\">\n{}\n</file>", r.uri, r.text))
198 - }
199 - // Describe binary resources rather than dropping them, so the model knows they exist.
200 - EmbeddedResourceResource::BlobResourceContents(b) => Some(format!(
201 - "[Attached binary resource: {} ({})]",
202 - b.uri,
203 - b.mime_type.as_deref().unwrap_or("unknown type")
204 - )),
205 - _ => None,
206 - }
207 -}
208 -
209 -/// How a resource link is shown to the model. Audio files are pointed at `analyze_audio`
210 -/// (and the separator) instead of being presented as text to read.
211 -fn link_text(link: &ResourceLink) -> String {
212 - match parse_file_link(&link.uri) {
213 - Some(FileLink { path, .. }) if is_audio(&path, link.mime_type.as_deref()) => format!(
214 - "[Referenced audio file: {}. Use analyze_audio on that path to inspect it.]",
215 - path.display()
216 - ),
217 - _ => format!("[Referenced: {}]", link.uri),
218 - }
219 -}
220 -
221 -/// File extensions the SplitFire app accepts, plus the other formats symphonia decodes.
222 -const AUDIO_EXTENSIONS: &[&str] = &[
223 - "wav", "wave", "mp3", "flac", "m4a", "aac", "mp4", "ogg", "oga", "opus", "aif", "aiff", "aifc",
224 - "caf", "webm", "mkv",
225 -];
226 -
227 -fn is_audio(path: &Path, mime: Option<&str>) -> bool {
228 - mime.is_some_and(|m| m.starts_with("audio/"))
229 - || path
230 - .extension()
231 - .and_then(|e| e.to_str())
232 - .is_some_and(|e| AUDIO_EXTENSIONS.contains(&e.to_ascii_lowercase().as_str()))
233 -}
234 -
235 -/// Largest file a `file://` resource link is inlined for; bigger ones stay as references.
236 -const MAX_LINKED_FILE_BYTES: usize = 256 * 1024;
237 -
238 -/// Replace `file://` resource links to text files with their embedded contents, so the model
239 -/// sees the file. Reads through the client when it supports `fs/read_text_file` (unsaved
240 -/// buffers), else from disk. A link to a selection (`#L10:20`) is inlined as just those
241 -/// lines. Audio files, links that can't be read and oversized files stay as references.
242 -pub async fn resolve_resource_links(
243 - blocks: &[ContentBlock],
244 - caps: &ClientCapabilities,
245 - connection: &ConnectionTo<Client>,
246 - session_id: &SessionId,
247 -) -> Vec<ContentBlock> {
248 - let mut out = Vec::with_capacity(blocks.len());
249 - for block in blocks {
250 - if let ContentBlock::ResourceLink(link) = block
251 - && let Some(FileLink { path, lines }) = parse_file_link(&link.uri)
252 - && !is_audio(&path, link.mime_type.as_deref())
253 - {
254 - let text = if caps.fs.read_text_file {
255 - connection
256 - .send_request(ReadTextFileRequest::new(session_id.clone(), path.clone()))
257 - .block_task()
258 - .await
259 - .map(|r| r.content)
260 - .map_err(|e| e.to_string())
261 - } else {
262 - tokio::fs::read_to_string(&path)
263 - .await
264 - .map_err(|e| e.to_string())
265 - };
266 - let text = match lines {
267 - Some(range) => text.map(|t| slice_lines(&t, range)),
268 - None => text,
269 - };
270 - match text {
271 - Ok(text) if lines.is_some() && text.is_empty() => {
272 - tracing::debug!("{} selects no lines", link.uri)
273 - }
274 - Ok(text) if text.len() <= MAX_LINKED_FILE_BYTES => {
275 - out.push(ContentBlock::Resource(EmbeddedResource::new(
276 - EmbeddedResourceResource::TextResourceContents(
277 - TextResourceContents::new(text, link.uri.clone())
278 - .mime_type(link.mime_type.clone()),
279 - ),
280 - )));
281 - continue;
282 - }
283 - Ok(_) => tracing::debug!("{} too large to inline", link.uri),
284 - Err(e) => tracing::debug!("could not read {}: {e}", link.uri),
285 - }
286 - }
287 - out.push(block.clone());
288 - }
289 - out
290 -}
291 -
292 -/// A `file://` resource link: the absolute path, and the lines it selects, if any.
293 -#[derive(Debug, PartialEq)]
294 -struct FileLink {
295 - path: PathBuf,
296 - /// 1-based, inclusive.
297 - lines: Option<(usize, usize)>,
298 -}
299 -
300 -/// Parse a `file://` URI, percent-decoding the path. Editors put more than the path in it:
301 -/// Zed links a selection as `file:///a.md?column=5#L10:20` and a symbol as
302 -/// `file:///a.rs?symbol=main#L3:9`. The query is dropped and the fragment read as a line
303 -/// range; a literal `?` or `#` in a file name arrives percent-encoded, so splitting first is
304 -/// safe.
305 -fn parse_file_link(uri: &str) -> Option<FileLink> {
306 - let rest = uri.strip_prefix("file://")?;
307 - let (rest, fragment) = match rest.split_once('#') {
308 - Some((rest, fragment)) => (rest, Some(fragment)),
309 - None => (rest, None),
310 - };
311 - let rest = rest.split_once('?').map_or(rest, |(rest, _)| rest);
312 - // Allow an authority of "" or "localhost"; reject other hosts.
313 - let path = if rest.starts_with('/') {
314 - rest
315 - } else {
316 - rest.strip_prefix("localhost")?
317 - };
318 - let bytes = path.as_bytes();
319 - let mut decoded = Vec::with_capacity(bytes.len());
320 - let mut i = 0;
321 - while i < bytes.len() {
322 - if bytes[i] == b'%'
323 - && i + 2 < bytes.len()
324 - && let Ok(b) = u8::from_str_radix(std::str::from_utf8(&bytes[i + 1..i + 3]).ok()?, 16)
325 - {
326 - decoded.push(b);
327 - i += 3;
328 - } else {
329 - decoded.push(bytes[i]);
330 - i += 1;
331 - }
332 - }
333 - let path = PathBuf::from(String::from_utf8(decoded).ok()?);
334 - path.is_absolute().then(|| FileLink {
335 - path,
336 - lines: fragment.and_then(parse_line_range),
337 - })
338 -}
339 -
340 -/// A `#L10:20` fragment as a 1-based inclusive range. Accepts `L10:20`, `L10-20`, `L10-L20`
341 -/// and a single line `L10`.
342 -fn parse_line_range(fragment: &str) -> Option<(usize, usize)> {
343 - let range = fragment.strip_prefix('L')?;
344 - let (start, end) = range
345 - .split_once(':')
346 - .or_else(|| range.split_once('-'))
347 - .unwrap_or((range, range));
348 - let end = end.strip_prefix('L').unwrap_or(end);
349 - let (start, end) = (start.parse::<usize>().ok()?, end.parse::<usize>().ok()?);
350 - (start >= 1 && end >= start).then_some((start, end))
351 -}
352 -
353 -/// Lines `start..=end` (1-based) of `text`, keeping their line endings.
354 -fn slice_lines(text: &str, (start, end): (usize, usize)) -> String {
355 - text.split_inclusive('\n')
356 - .skip(start - 1)
357 - .take(end - start + 1)
358 - .collect()
359 -}
360 -
361 -/// Flatten text-like prompt blocks into one string (used for titles and slash commands).
362 -pub fn prompt_to_text(blocks: &[ContentBlock]) -> String {
363 - let mut parts = Vec::new();
364 - for block in blocks {
365 - match block {
366 - ContentBlock::Text(t) => parts.push(t.text.clone()),
367 - ContentBlock::ResourceLink(link) => parts.push(link_text(link)),
368 - ContentBlock::Resource(res) => parts.extend(embedded_text(res)),
369 - _ => {}
370 - }
371 - }
372 - parts.join("\n\n")
373 -}
374 -
375 -fn audio_extension(mime: &str) -> &'static str {
376 - match mime.split(';').next().unwrap_or("").trim() {
377 - "audio/wav" | "audio/x-wav" | "audio/wave" | "audio/vnd.wave" => "wav",
378 - "audio/mpeg" | "audio/mp3" => "mp3",
379 - "audio/flac" | "audio/x-flac" => "flac",
380 - "audio/ogg" | "audio/vorbis" | "application/ogg" => "ogg",
381 - "audio/mp4" | "audio/m4a" | "audio/x-m4a" | "audio/aac" => "m4a",
382 - "audio/aiff" | "audio/x-aiff" => "aiff",
383 - "audio/webm" | "video/webm" => "webm",
384 - _ => "audio",
385 - }
386 -}
387 -
388 -/// Write each audio block under `dir`; returns the saved path per audio block, in order.
389 -pub async fn save_audio_blocks(blocks: &[ContentBlock], dir: &Path) -> Result<Vec<PathBuf>> {
390 - let mut saved = Vec::new();
391 - for block in blocks {
392 - let ContentBlock::Audio(audio) = block else {
393 - continue;
394 - };
395 - let bytes = base64::engine::general_purpose::STANDARD
396 - .decode(audio.data.trim())
397 - .context("audio block is not valid base64")?;
398 - tokio::fs::create_dir_all(dir).await?;
399 - let path = dir.join(format!(
400 - "audio-{}.{}",
401 - uuid::Uuid::new_v4().simple(),
402 - audio_extension(&audio.mime_type)
403 - ));
404 - tokio::fs::write(&path, bytes).await?;
405 - saved.push(path);
406 - }
407 - Ok(saved)
408 -}
409 -
410 -/// Convert prompt blocks to an OpenAI `content` value: a string for text only, or parts when
411 -/// images are present. Audio is referenced by the path it was saved to; the model reads it
412 -/// with `analyze_audio`. `slash` replaces the first text block (an expanded slash command).
413 -pub fn prompt_to_content(
414 - blocks: &[ContentBlock],
415 - audio_paths: &[PathBuf],
416 - slash: Option<&str>,
417 -) -> Value {
418 - let mut audio = audio_paths.iter();
419 - let mut parts: Vec<Value> = Vec::new();
420 - let mut slash = slash;
421 - let mut has_image = false;
422 - let text = |t: String| json!({ "type": "text", "text": t });
423 - for block in blocks {
424 - match block {
425 - ContentBlock::Text(t) => match slash.take() {
426 - Some(expanded) => parts.push(text(expanded.to_string())),
427 - None => parts.push(text(t.text.clone())),
428 - },
429 - ContentBlock::ResourceLink(link) => parts.push(text(link_text(link))),
430 - ContentBlock::Resource(res) => {
431 - if let Some(t) = embedded_text(res) {
432 - parts.push(text(t));
433 - }
434 - }
435 - ContentBlock::Image(img) => {
436 - has_image = true;
437 - let url = format!("data:{};base64,{}", img.mime_type, img.data);
438 - parts.push(json!({ "type": "image_url", "image_url": { "url": url } }));
439 - }
440 - ContentBlock::Audio(a) => match audio.next() {
441 - Some(path) => parts.push(text(format!(
442 - "[Attached audio ({}) saved at {}. Use analyze_audio on that path to inspect it.]",
443 - a.mime_type,
444 - path.display()
445 - ))),
446 - None => parts.push(text("[Attached audio could not be saved]".into())),
447 - },
448 - _ => {}
449 - }
450 - }
451 - if has_image {
452 - return Value::Array(parts);
453 - }
454 - let joined = parts
455 - .iter()
456 - .filter_map(|p| p.get("text").and_then(Value::as_str))
457 - .collect::<Vec<_>>()
458 - .join("\n\n");
459 - Value::String(joined)
460 -}
461 -
187 #[cfg(test)]
188 mod tests {
189 use super::*;
190
466 - fn link(uri: &str) -> Option<FileLink> {
467 - parse_file_link(uri)
468 - }
469 -
470 - #[test]
471 - fn plain_paths_decode_and_must_be_absolute() {
472 - assert_eq!(
473 - link("file:///a%20b/c.md"),
474 - Some(FileLink {
475 - path: "/a b/c.md".into(),
476 - lines: None
477 - })
478 - );
479 - assert_eq!(
480 - link("file://localhost/x.md").unwrap().path,
481 - Path::new("/x.md")
482 - );
483 - assert_eq!(link("file://other-host/x.md"), None);
484 - assert_eq!(link("https://example.com/x.md"), None);
485 - }
486 -
487 - #[test]
488 - fn zed_selection_and_symbol_links() {
489 - let sel = link("file:///notes/song.md?column=5#L10:20").unwrap();
490 - assert_eq!(sel.path, Path::new("/notes/song.md"));
491 - assert_eq!(sel.lines, Some((10, 20)));
492 - assert_eq!(
493 - link("file:///a.rs?symbol=main#L3:9").unwrap().lines,
494 - Some((3, 9))
495 - );
496 - }
497 -
498 - #[test]
499 - fn line_range_forms() {
500 - assert_eq!(parse_line_range("L10:20"), Some((10, 20)));
501 - assert_eq!(parse_line_range("L10-20"), Some((10, 20)));
502 - assert_eq!(parse_line_range("L10-L20"), Some((10, 20)));
503 - assert_eq!(parse_line_range("L7"), Some((7, 7)));
504 - assert_eq!(parse_line_range("L0"), None);
505 - assert_eq!(parse_line_range("L5:2"), None);
506 - assert_eq!(parse_line_range("x"), None);
507 - }
508 -
509 - #[test]
510 - fn slices_inclusive_lines() {
511 - assert_eq!(slice_lines("a\nb\nc\nd\n", (2, 3)), "b\nc\n");
512 - assert_eq!(slice_lines("a\nb", (2, 9)), "b");
513 - assert_eq!(slice_lines("a\n", (5, 6)), "");
514 - }
515 -
516 - #[test]
517 - fn audio_links_point_at_the_analyzer() {
518 - let wav = ResourceLink::new("mix.wav", "file:///music/mix%20v2.WAV");
519 - assert_eq!(
520 - link_text(&wav),
521 - "[Referenced audio file: /music/mix v2.WAV. Use analyze_audio on that path to inspect it.]"
522 - );
523 - let by_mime =
524 - ResourceLink::new("take", "file:///t/take").mime_type("audio/flac".to_string());
525 - assert!(link_text(&by_mime).contains("analyze_audio"));
526 - let md = ResourceLink::new("notes", "file:///t/notes.md");
527 - assert_eq!(link_text(&md), "[Referenced: file:///t/notes.md]");
528 - }
529 - use agent_client_protocol::schema::v1::{AudioContent, ImageContent, TextContent};
530 -
191 #[test]
192 fn system_prompt_mentions_stems_only_when_connected() {
193 let cwd = Path::new("/work");
@@ -595,43 +255,4 @@ mod tests {
255 assert!(expand_slash("/unknown x", false).is_none());
256 assert!(expand_slash("hello /analyze", false).is_none());
257 }
598 -
599 - #[tokio::test]
600 - async fn audio_is_saved_and_referenced() {
601 - let dir = std::env::temp_dir().join(format!("sf-prompt-{}", uuid::Uuid::new_v4()));
602 - let data = base64::engine::general_purpose::STANDARD.encode(b"RIFFfake");
603 - let blocks = vec![
604 - ContentBlock::Text(TextContent::new("what key?")),
605 - ContentBlock::Audio(AudioContent::new(data, "audio/wav")),
606 - ];
607 - let saved = save_audio_blocks(&blocks, &dir).await.unwrap();
608 - assert_eq!(saved.len(), 1);
609 - assert_eq!(saved[0].extension().unwrap(), "wav");
610 - assert_eq!(std::fs::read(&saved[0]).unwrap(), b"RIFFfake");
611 - let content = prompt_to_content(&blocks, &saved, None);
612 - let s = content.as_str().unwrap();
613 - assert!(s.contains("what key?") && s.contains(&saved[0].display().to_string()));
614 - std::fs::remove_dir_all(dir).ok();
615 -
616 - let bad = vec![ContentBlock::Audio(AudioContent::new("!!!", "audio/wav"))];
617 - assert!(
618 - save_audio_blocks(&bad, &std::env::temp_dir())
619 - .await
620 - .is_err()
621 - );
622 - }
623 -
624 - #[test]
625 - fn images_use_parts_and_slash_replaces_text() {
626 - let blocks = vec![
627 - ContentBlock::Text(TextContent::new("/chords Am F")),
628 - ContentBlock::Image(ImageContent::new("AAAA", "image/png")),
629 - ];
630 - let v = prompt_to_content(&blocks, &[], Some("EXPANDED"));
631 - let parts = v.as_array().unwrap();
632 - assert_eq!(parts[0]["text"], "EXPANDED");
633 - assert_eq!(parts[1]["type"], "image_url");
634 - let plain = prompt_to_content(&blocks[..1], &[], None);
635 - assert_eq!(plain, json!("/chords Am F"));
636 - }
258 }
src/session_store.rs deleted
-174
@@ -1,174 +0,0 @@
1 -//! Durable sessions: one JSON file per session under `<data_dir>/sessions/`.
2 -//!
3 -//! The history stored here is the OpenAI-style message list without the system prompt
4 -//! (that is rebuilt every turn). Tool results are `role: "tool"` messages, which is what
5 -//! lets `session/load` replay tool calls with their output.
6 -
7 -use std::path::{Path, PathBuf};
8 -
9 -use anyhow::{Context, Result, bail};
10 -use serde::{Deserialize, Serialize};
11 -use serde_json::Value;
12 -
13 -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
14 -pub struct StoredSession {
15 - pub cwd: PathBuf,
16 - #[serde(default)]
17 - pub roots: Vec<PathBuf>,
18 - pub model: String,
19 - #[serde(default)]
20 - pub title: Option<String>,
21 - pub updated_at: u64,
22 - #[serde(default)]
23 - pub messages: Vec<Value>,
24 -}
25 -
26 -#[derive(Debug, Clone)]
27 -pub struct SessionStore {
28 - dir: PathBuf,
29 -}
30 -
31 -/// Session ids end up in file names: allow only characters that cannot escape the directory.
32 -pub fn sanitize_id(id: &str) -> Result<&str> {
33 - let ok = !id.is_empty()
34 - && id.len() <= 128
35 - && id
36 - .chars()
37 - .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_');
38 - if !ok {
39 - bail!("invalid session id");
40 - }
41 - Ok(id)
42 -}
43 -
44 -impl SessionStore {
45 - pub fn new(data_dir: &Path) -> Self {
46 - Self {
47 - dir: data_dir.join("sessions"),
48 - }
49 - }
50 -
51 - fn path(&self, id: &str) -> Result<PathBuf> {
52 - Ok(self.dir.join(format!("{}.json", sanitize_id(id)?)))
53 - }
54 -
55 - /// Atomic write: temp file in the same directory, then rename.
56 - pub fn save(&self, id: &str, session: &StoredSession) -> Result<()> {
57 - let path = self.path(id)?;
58 - std::fs::create_dir_all(&self.dir)
59 - .with_context(|| format!("creating {}", self.dir.display()))?;
60 - let tmp = self.dir.join(format!(".{id}.{}.tmp", uuid::Uuid::new_v4()));
61 - std::fs::write(&tmp, serde_json::to_vec(session)?)?;
62 - std::fs::rename(&tmp, &path).inspect_err(|_| {
63 - let _ = std::fs::remove_file(&tmp);
64 - })?;
65 - Ok(())
66 - }
67 -
68 - /// `Ok(None)` when no such session exists (or the id is not a valid one).
69 - pub fn load(&self, id: &str) -> Result<Option<StoredSession>> {
70 - let Ok(path) = self.path(id) else {
71 - return Ok(None);
72 - };
73 - match std::fs::read(&path) {
74 - Ok(bytes) => Ok(Some(
75 - serde_json::from_slice(&bytes)
76 - .with_context(|| format!("parsing {}", path.display()))?,
77 - )),
78 - Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
79 - Err(e) => Err(e).with_context(|| format!("reading {}", path.display())),
80 - }
81 - }
82 -
83 - /// All stored sessions, newest first. Unreadable files are skipped.
84 - pub fn list(&self) -> Vec<(String, StoredSession)> {
85 - let Ok(rd) = std::fs::read_dir(&self.dir) else {
86 - return Vec::new();
87 - };
88 - let mut out: Vec<(String, StoredSession)> = rd
89 - .filter_map(|e| {
90 - let path = e.ok()?.path();
91 - if path.extension()? != "json" {
92 - return None;
93 - }
94 - let id = path.file_stem()?.to_str()?.to_string();
95 - let session = serde_json::from_slice(&std::fs::read(&path).ok()?).ok()?;
96 - Some((id, session))
97 - })
98 - .collect();
99 - out.sort_by(|a, b| b.1.updated_at.cmp(&a.1.updated_at).then(a.0.cmp(&b.0)));
100 - out
101 - }
102 -
103 - /// Returns whether a file was removed.
104 - pub fn delete(&self, id: &str) -> Result<bool> {
105 - let Ok(path) = self.path(id) else {
106 - return Ok(false);
107 - };
108 - match std::fs::remove_file(&path) {
109 - Ok(()) => Ok(true),
110 - Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
111 - Err(e) => Err(e).with_context(|| format!("deleting {}", path.display())),
112 - }
113 - }
114 -}
115 -
116 -#[cfg(test)]
117 -mod tests {
118 - use super::*;
119 - use serde_json::json;
120 -
121 - fn tmp() -> PathBuf {
122 - let d = std::env::temp_dir().join(format!("sf-store-{}", uuid::Uuid::new_v4()));
123 - std::fs::create_dir_all(&d).unwrap();
124 - d
125 - }
126 -
127 - fn sample(updated_at: u64) -> StoredSession {
128 - StoredSession {
129 - cwd: "/work".into(),
130 - roots: vec!["/other".into()],
131 - model: "m".into(),
132 - title: Some("hello".into()),
133 - updated_at,
134 - messages: vec![
135 - json!({"role": "user", "content": "hi"}),
136 - json!({"role": "tool", "tool_call_id": "c1", "content": "out"}),
137 - ],
138 - }
139 - }
140 -
141 - #[test]
142 - fn round_trip_list_delete() {
143 - let dir = tmp();
144 - let store = SessionStore::new(&dir);
145 - assert!(store.load("a").unwrap().is_none());
146 - store.save("a", &sample(10)).unwrap();
147 - store.save("b", &sample(20)).unwrap();
148 - assert_eq!(store.load("a").unwrap().unwrap(), sample(10));
149 - let ids: Vec<_> = store.list().into_iter().map(|(id, _)| id).collect();
150 - assert_eq!(ids, ["b", "a"]);
151 - assert!(store.delete("a").unwrap());
152 - assert!(!store.delete("a").unwrap());
153 - assert!(store.load("a").unwrap().is_none());
154 - // No temp files left behind.
155 - let leftovers = std::fs::read_dir(dir.join("sessions"))
156 - .unwrap()
157 - .filter(|e| e.as_ref().unwrap().path().extension().unwrap() != "json")
158 - .count();
159 - assert_eq!(leftovers, 0);
160 - std::fs::remove_dir_all(dir).ok();
161 - }
162 -
163 - #[test]
164 - fn rejects_path_traversal_ids() {
165 - let dir = tmp();
166 - let store = SessionStore::new(&dir);
167 - for bad in ["../x", "a/b", "", "a.b", "..", "a\\b"] {
168 - assert!(store.save(bad, &sample(1)).is_err(), "{bad}");
169 - assert!(store.load(bad).unwrap().is_none());
170 - assert!(!store.delete(bad).unwrap());
171 - }
172 - std::fs::remove_dir_all(dir).ok();
173 - }
174 -}
src/tools/mod.rs
+55 -395
@@ -1,417 +1,77 @@
1 -//! Tools exposed to the model: workspace files/shell, the music theory engine, audio
2 -//! analysis, and tools from MCP servers (client-supplied plus the optional demucs server).
1 +//! SplitFire's tools: the music theory engine, audio analysis and the separation library.
2 +//! Workspace files and shell, MCP and approval come from `ed-acp`.
3
4 pub mod audio;
5 pub mod library;
6 pub mod theory;
7 -mod workspace;
7
9 -use std::collections::HashSet;
10 -use std::path::{Component, Path, PathBuf};
11 -use std::sync::atomic::{AtomicBool, Ordering};
12 -use std::sync::{Arc, Mutex};
8 +use agent_client_protocol::schema::v1::{ToolCallLocation, ToolKind};
9 +use anyhow::{Result, anyhow};
10 +use async_trait::async_trait;
11 +use serde_json::Value;
12
14 -use agent_client_protocol::schema::v1::{
15 - ClientCapabilities, PermissionOption, PermissionOptionKind, RequestPermissionOutcome,
16 - RequestPermissionRequest, SessionId, SessionNotification, SessionUpdate, ToolCallContent,
17 - ToolCallLocation, ToolCallUpdate, ToolCallUpdateFields, ToolKind,
13 +pub use ed_acp::{
14 + DescribeCtx, ToolCtx, ToolOutcome, Toolset, absolutize, function_def, resolve_in,
15 };
19 -use agent_client_protocol::{Client, ConnectionTo};
20 -use anyhow::{Result, anyhow, bail};
21 -use serde_json::{Value, json};
22 -use tokio_util::sync::CancellationToken;
16
24 -use crate::mcp::{self, McpToolset};
17 +/// The music tools as one [`Toolset`]: six `theory_*` tools, `analyze_audio` and
18 +/// `list_separations`. All are read-only and never ask for approval.
19 +#[derive(Debug, Default, Clone, Copy)]
20 +pub struct MusicTools;
21
26 -pub(crate) const MAX_OUTPUT_BYTES: usize = 32 * 1024;
27 -
28 -/// Everything a tool needs to act on behalf of one session.
29 -pub struct ToolCtx {
30 - pub connection: ConnectionTo<Client>,
31 - pub session_id: SessionId,
32 - /// Primary working directory: the base for relative paths and `run_command`.
33 - pub cwd: PathBuf,
34 - /// Additional workspace roots the model may address with the `root` parameter.
35 - pub roots: Vec<PathBuf>,
36 - pub caps: ClientCapabilities,
37 - pub cancel: CancellationToken,
38 - /// Approve tool calls without asking: `--yolo`, or the session's `auto_approve` option.
39 - pub auto_approve: Arc<AtomicBool>,
40 - /// Tool names the user chose "always allow" for in this session.
41 - pub always_allowed: Arc<Mutex<HashSet<String>>>,
42 - /// Tool names the user chose "always reject" for in this session.
43 - pub always_rejected: Arc<Mutex<HashSet<String>>>,
44 - /// Tools from MCP servers available to this session.
45 - pub mcp: Arc<McpToolset>,
46 -}
22 +#[async_trait]
23 +impl Toolset for MusicTools {
24 + fn definitions(&self) -> Vec<Value> {
25 + let mut defs = theory::definitions();
26 + defs.push(audio::definition());
27 + defs.push(library::definition());
28 + defs
29 + }
30
48 -pub struct ToolOutcome {
49 - /// Text returned to the model.
50 - pub text: String,
51 - /// Rich content shown to the user in the client.
52 - pub content: Vec<ToolCallContent>,
53 - pub failed: bool,
54 - /// Files the tool produced or touched, surfaced as tool-call locations.
55 - pub locations: Vec<PathBuf>,
56 -}
31 + fn handles(&self, name: &str) -> bool {
32 + name.starts_with("theory_") || name == audio::NAME || name == library::NAME
33 + }
34
58 -impl ToolOutcome {
59 - pub fn ok(text: impl Into<String>) -> Self {
60 - let text = text.into();
61 - Self {
62 - content: vec![ToolCallContent::from(text.clone())],
63 - text,
64 - failed: false,
65 - locations: Vec::new(),
35 + fn describe(
36 + &self,
37 + name: &str,
38 + args: &Value,
39 + ctx: &DescribeCtx<'_>,
40 + ) -> (String, ToolKind, Vec<ToolCallLocation>) {
41 + if name.starts_with("theory_") {
42 + return (theory::title(name, args), ToolKind::Think, vec![]);
43 }
67 - }
68 - pub fn err(text: impl Into<String>) -> Self {
69 - let text = text.into();
70 - Self {
71 - content: vec![ToolCallContent::from(text.clone())],
72 - text,
73 - failed: true,
74 - locations: Vec::new(),
44 + if name == library::NAME {
45 + return ("List separated songs".to_string(), ToolKind::Search, vec![]);
46 }
76 - }
77 -}
78 -
79 -pub(crate) fn function_def(name: &str, desc: &str, params: Value) -> Value {
80 - json!({ "type": "function", "function": { "name": name, "description": desc, "parameters": params } })
81 -}
82 -
83 -/// All tool definitions for one session, in OpenAI `tools` format.
84 -pub fn definitions(mcp: &McpToolset) -> Value {
85 - let mut defs = workspace::definitions();
86 - defs.extend(theory::definitions());
87 - defs.push(audio::definition());
88 - defs.push(library::definition());
89 - defs.extend(mcp.definitions());
90 - Value::Array(defs)
91 -}
92 -
93 -/// Title, kind and affected locations shown in the client when the call starts.
94 -pub fn describe(
95 - cwd: &Path,
96 - roots: &[PathBuf],
97 - name: &str,
98 - args: &Value,
99 -) -> (String, ToolKind, Vec<ToolCallLocation>) {
100 - let str_arg = |k: &str| args.get(k).and_then(Value::as_str);
101 - if let Some((server, tool)) = mcp::split_qualified(name) {
102 - let loc = ["input", "path", "file"]
103 - .iter()
104 - .filter_map(|k| str_arg(k))
105 - .map(PathBuf::from)
106 - .filter(|p| p.is_absolute())
107 - .map(|p| ToolCallLocation::new(absolutize(&p)))
108 - .collect();
109 - return (format!("{server}: {tool}"), ToolKind::Other, loc);
110 - }
111 - if name.starts_with("theory_") {
112 - return (theory::title(name, args), ToolKind::Think, vec![]);
113 - }
114 - if name == library::NAME {
115 - return ("List separated songs".to_string(), ToolKind::Search, vec![]);
116 - }
117 - if name == audio::NAME {
118 - let file = str_arg("path").unwrap_or("audio file");
119 - let loc = str_arg("path")
120 - .map(|p| absolutize(&resolve_in(cwd, p)))
47 + let path = args.get("path").and_then(Value::as_str);
48 + let loc = path
49 + .map(|p| absolutize(&resolve_in(ctx.cwd, p)))
50 .map(ToolCallLocation::new)
51 .into_iter()
52 .collect();
124 - return (format!("Analyze {file}"), ToolKind::Read, loc);
53 + (
54 + format!("Analyze {}", path.unwrap_or("audio file")),
55 + ToolKind::Read,
56 + loc,
57 + )
58 }
126 - let base = base_for(cwd, roots, str_arg("root")).unwrap_or_else(|_| cwd.to_path_buf());
127 - // ACP v1: ToolCallLocation.path must be an absolute path.
128 - let path = str_arg("path").map(|p| absolutize(&resolve_in(&base, p)));
129 - let loc = path
130 - .clone()
131 - .map(ToolCallLocation::new)
132 - .into_iter()
133 - .collect();
134 - let shown = path
135 - .as_deref()
136 - .map_or_else(|| base.display().to_string(), |p| p.display().to_string());
137 - match name {
138 - "read_file" => (format!("Read {shown}"), ToolKind::Read, loc),
139 - "write_file" => (format!("Write {shown}"), ToolKind::Edit, loc),
140 - "edit_file" => (format!("Edit {shown}"), ToolKind::Edit, loc),
141 - "list_directory" => (format!("List {shown}"), ToolKind::Search, loc),
142 - "run_command" => {
143 - let cmd = str_arg("command").unwrap_or("");
144 - let where_ = if base == cwd {
145 - String::new()
146 - } else {
147 - format!(" in {}", base.display())
148 - };
149 - (format!("`{cmd}`{where_}"), ToolKind::Execute, vec![])
150 - }
151 - _ => (name.to_string(), ToolKind::Other, vec![]),
152 - }
153 -}
154 -
155 -pub async fn execute(ctx: &ToolCtx, tool_call_id: &str, name: &str, args: Value) -> ToolOutcome {
156 - let result = if workspace::NAMES.contains(&name) {
157 - workspace::execute(ctx, tool_call_id, name, args).await
158 - } else if name.starts_with("theory_") {
159 - theory::execute(name, args)
160 - } else if name == audio::NAME {
161 - audio::execute(ctx, args).await
162 - } else if name == library::NAME {
163 - library::execute(ctx, args).await
164 - } else if mcp::split_qualified(name).is_some() {
165 - execute_mcp(ctx, tool_call_id, name, args).await
166 - } else {
167 - Err(anyhow!("unknown tool `{name}`"))
168 - };
169 - result.unwrap_or_else(|e| ToolOutcome::err(format!("Error: {e:#}")))
170 -}
171 -
172 -async fn execute_mcp(
173 - ctx: &ToolCtx,
174 - tool_call_id: &str,
175 - name: &str,
176 - args: Value,
177 -) -> Result<ToolOutcome> {
178 - if !ctx.mcp.has(name) {
179 - bail!("unknown tool `{name}`");
180 - }
181 - if !ctx.mcp.is_read_only(name) {
182 - let preview = ctx.mcp.permission_preview(name, &args).await;
183 - if !ctx.permit(tool_call_id, name, preview).await? {
184 - return Ok(ToolOutcome::err("User rejected running this tool."));
185 - }
186 - }
187 - ctx.mcp.call(ctx, tool_call_id, name, args).await
188 -}
59
190 -/// The base directory for a tool call's `root` parameter: a workspace root or any directory
191 -/// beneath one (so the model can pass e.g. `root/crate/src`). Paths that escape a root via
192 -/// `..` are rejected.
193 -pub(crate) fn base_for(cwd: &Path, roots: &[PathBuf], root: Option<&str>) -> Result<PathBuf> {
194 - let Some(root) = root else {
195 - return Ok(cwd.to_path_buf());
196 - };
197 - // Lexical normalization resolves `..`, so `root/../elsewhere` can't pass the prefix check.
198 - let candidate = normalize(Path::new(root));
199 - let all_roots = all_roots(cwd, roots);
200 - for ws_root in &all_roots {
201 - if candidate == *ws_root {
202 - return Ok(ws_root.clone());
203 - }
204 - // Allow any directory below a root (e.g. a crate dir like `root/crate/src`),
205 - // as long as it doesn't escape the root via `..`.
206 - if let Ok(rest) = candidate.strip_prefix(ws_root)
207 - && rest.components().all(|c| matches!(c, Component::Normal(_)))
208 - {
209 - return Ok(candidate);
210 - }
211 - }
212 - let known = all_roots
213 - .iter()
214 - .map(|r| r.display().to_string())
215 - .collect::<Vec<_>>()
216 - .join(", ");
217 - bail!(
218 - "unknown root `{}` (workspace roots: {known})",
219 - candidate.display()
220 - )
221 -}
222 -
223 -fn all_roots(cwd: &Path, roots: &[PathBuf]) -> Vec<PathBuf> {
224 - std::iter::once(cwd.to_path_buf())
225 - .chain(roots.iter().filter(|r| *r != cwd).cloned())
226 - .collect()
227 -}
228 -
229 -impl ToolCtx {
230 - pub(crate) fn base_for(&self, root: Option<&str>) -> Result<PathBuf> {
231 - base_for(&self.cwd, &self.roots, root)
232 - }
233 -
234 - pub(crate) fn resolve_with(&self, root: Option<&str>, path: &str) -> Result<PathBuf> {
235 - Ok(resolve_in(&self.base_for(root)?, path))
236 - }
237 -
238 - /// Ask the user for permission. Returns `Ok(false)` if rejected or cancelled.
239 - pub(crate) async fn permit(
60 + async fn execute(
61 &self,
241 - id: &str,
242 - tool: &str,
243 - content: Vec<ToolCallContent>,
244 - ) -> Result<bool> {
245 - if self.always_rejected.lock().unwrap().contains(tool) {
246 - return Ok(false);
247 - }
248 - if self.auto_approve.load(Ordering::Relaxed)
249 - || self.always_allowed.lock().unwrap().contains(tool)
250 - {
251 - return Ok(true);
252 - }
253 - let mut fields = ToolCallUpdateFields::new();
254 - if !content.is_empty() {
255 - fields = fields.content(content);
256 - }
257 - let req = RequestPermissionRequest::new(
258 - self.session_id.clone(),
259 - ToolCallUpdate::new(id.to_string(), fields),
260 - vec![
261 - PermissionOption::new("allow_once", "Allow", PermissionOptionKind::AllowOnce),
262 - PermissionOption::new(
263 - "allow_always",
264 - "Always allow",
265 - PermissionOptionKind::AllowAlways,
266 - ),
267 - PermissionOption::new("reject_once", "Reject", PermissionOptionKind::RejectOnce),
268 - PermissionOption::new(
269 - "reject_always",
270 - "Always reject",
271 - PermissionOptionKind::RejectAlways,
272 - ),
273 - ],
274 - );
275 - let resp = tokio::select! {
276 - r = self.connection.send_request(req).block_task() => r?,
277 - () = self.cancel.cancelled() => return Ok(false),
278 - };
279 - Ok(match resp.outcome {
280 - RequestPermissionOutcome::Selected(sel) => match &*sel.option_id.0 {
281 - "allow_always" => {
282 - self.always_allowed.lock().unwrap().insert(tool.to_string());
283 - true
284 - }
285 - "allow_once" => true,
286 - "reject_always" => {
287 - self.always_rejected
288 - .lock()
289 - .unwrap()
290 - .insert(tool.to_string());
291 - false
292 - }
293 - _ => false,
294 - },
295 - _ => false,
296 - })
297 - }
298 -
299 - pub fn update(&self, id: &str, fields: ToolCallUpdateFields) -> Result<()> {
300 - let update = SessionUpdate::ToolCallUpdate(ToolCallUpdate::new(id.to_string(), fields));
301 - self.connection
302 - .send_notification(SessionNotification::new(self.session_id.clone(), update))?;
303 - Ok(())
304 - }
305 -}
306 -
307 -/// Join `path` onto `base`, keeping absolute paths as-is.
308 -pub(crate) fn resolve_in(base: &Path, path: &str) -> PathBuf {
309 - let p = Path::new(path);
310 - if p.is_absolute() {
311 - p.to_path_buf()
312 - } else {
313 - base.join(p)
314 - }
315 -}
316 -
317 -pub(crate) fn truncate(mut s: String) -> String {
318 - if s.len() > MAX_OUTPUT_BYTES {
319 - let mut cut = MAX_OUTPUT_BYTES;
320 - while !s.is_char_boundary(cut) {
321 - cut -= 1;
62 + ctx: &ToolCtx,
63 + _tool_call_id: &str,
64 + name: &str,
65 + args: Value,
66 + ) -> Result<ToolOutcome> {
67 + if name.starts_with("theory_") {
68 + theory::execute(name, args)
69 + } else if name == audio::NAME {
70 + audio::execute(ctx, args).await
71 + } else if name == library::NAME {
72 + library::execute(ctx, args).await
73 + } else {
74 + Err(anyhow!("unknown tool `{name}`"))
75 }
323 - s.truncate(cut);
324 - s.push_str("\n[truncated]");
325 - }
326 - s
327 -}
328 -
329 -/// Make `path` absolute and free of `.`/`..` components, as the ACP v1 spec requires
330 -/// absolute paths everywhere. Existing paths are canonicalized (resolving symlinks
331 -/// and giving the client the true on-disk location); paths that don't exist yet
332 -/// are normalized lexically with the existing parent canonicalized when possible.
333 -pub(crate) fn absolutize(path: &Path) -> PathBuf {
334 - if let Ok(c) = std::fs::canonicalize(path) {
335 - return c;
336 - }
337 - // Not on disk (yet): normalize lexically, anchoring at the canonicalized parent.
338 - let parent = path.parent().filter(|p| !p.as_os_str().is_empty());
339 - let base = parent.and_then(|p| std::fs::canonicalize(p).ok());
340 - let file_name = path.file_name().map(|f| f.to_os_string());
341 - match (base, file_name) {
342 - (Some(mut b), Some(f)) => {
343 - b.push(f);
344 - b
345 - }
346 - _ => normalize(path),
347 - }
348 -}
349 -
350 -/// Lexically normalize a path: resolve `.` and `..` without touching the filesystem.
351 -/// Absolute paths stay absolute; relative paths are made absolute against the cwd.
352 -pub(crate) fn normalize(path: &Path) -> PathBuf {
353 - let absolute = if path.is_absolute() {
354 - path.to_path_buf()
355 - } else {
356 - std::env::current_dir()
357 - .unwrap_or_else(|_| PathBuf::from("."))
358 - .join(path)
359 - };
360 - let mut out: Vec<std::ffi::OsString> = Vec::new();
361 - for c in absolute.components() {
362 - match c {
363 - Component::CurDir => {}
364 - Component::ParentDir => {
365 - out.pop();
366 - }
367 - other => out.push(other.as_os_str().to_os_string()),
368 - }
369 - }
370 - out.into_iter().collect()
371 -}
372 -
373 -#[cfg(test)]
374 -mod fs_path_tests {
375 - use super::*;
376 -
377 - #[test]
378 - fn normalize_resolves_dot_dot_lexically() {
379 - assert_eq!(normalize(Path::new("/a/b/../c")), PathBuf::from("/a/c"));
380 - assert_eq!(normalize(Path::new("/a/./b")), PathBuf::from("/a/b"));
381 - assert_eq!(normalize(Path::new("/a/b/..")), PathBuf::from("/a"));
382 - }
383 -
384 - #[test]
385 - fn normalize_makes_relative_paths_absolute() {
386 - let n = normalize(Path::new("src/../README.md"));
387 - assert!(n.is_absolute());
388 - assert!(n.ends_with("README.md"));
389 - assert!(!n.components().any(|c| c == Component::ParentDir));
390 - }
391 -
392 - #[test]
393 - fn absolutize_canonicalizes_existing_paths() {
394 - let n = absolutize(&std::env::temp_dir());
395 - assert!(n.is_absolute());
396 - assert!(!n.components().any(|c| c == Component::ParentDir));
397 - }
398 -
399 - #[test]
400 - fn absolutize_normalizes_missing_paths() {
401 - let n = absolutize(Path::new("/definitely/does/not/../exist.txt"));
402 - assert_eq!(n, PathBuf::from("/definitely/does/exist.txt"));
403 - }
404 -
405 - #[test]
406 - fn base_for_rejects_escapes_and_accepts_subdirs() {
407 - let cwd = PathBuf::from("/ws/a");
408 - let roots = vec![PathBuf::from("/ws/b")];
409 - assert_eq!(base_for(&cwd, &roots, None).unwrap(), cwd);
410 - assert_eq!(
411 - base_for(&cwd, &roots, Some("/ws/b/sub/../x")).unwrap(),
412 - PathBuf::from("/ws/b/x")
413 - );
414 - assert!(base_for(&cwd, &roots, Some("/ws/b/../c")).is_err());
415 - assert!(base_for(&cwd, &roots, Some("/etc")).is_err());
76 }
77 }
src/tools/workspace.rs deleted
-400
@@ -1,400 +0,0 @@
1 -//! Workspace tools: files and shell, routed through the ACP client when it advertises
2 -//! support and falling back to the local filesystem/process otherwise.
3 -
4 -use std::path::Path;
5 -use std::time::Duration;
6 -
7 -use agent_client_protocol::ErrorCode;
8 -use agent_client_protocol::schema::v1::{
9 - CreateTerminalRequest, Diff, KillTerminalRequest, ReadTextFileRequest, ReleaseTerminalRequest,
10 - Terminal, TerminalOutputRequest, ToolCallContent, ToolCallUpdateFields,
11 - WaitForTerminalExitRequest, WriteTextFileRequest,
12 -};
13 -use anyhow::{Context, Result, anyhow, bail};
14 -use serde::Deserialize;
15 -use serde_json::{Value, json};
16 -
17 -use super::{MAX_OUTPUT_BYTES, ToolCtx, ToolOutcome, absolutize, resolve_in, truncate};
18 -
19 -const COMMAND_TIMEOUT: Duration = Duration::from_secs(120);
20 -
21 -pub fn definitions() -> Vec<Value> {
22 - let f = |name: &str, desc: &str, params: Value| super::function_def(name, desc, params);
23 - let root_prop = || {
24 - json!({
25 - "type": "string",
26 - "description": "Workspace root to act in: absolute path of a session root or a directory beneath one. Defaults to the primary working directory."
27 - })
28 - };
29 - vec![
30 - f(
31 - "read_file",
32 - "Read a text file. Paths may be relative to the working directory.",
33 - json!({
34 - "type": "object",
35 - "properties": {
36 - "path": { "type": "string" },
37 - "line": { "type": "integer", "description": "1-based line to start from" },
38 - "limit": { "type": "integer", "description": "Max lines to read" },
39 - "root": root_prop()
40 - },
41 - "required": ["path"]
42 - }),
43 - ),
44 - f(
45 - "write_file",
46 - "Create or overwrite a file with the given content.",
47 - json!({
48 - "type": "object",
49 - "properties": { "path": { "type": "string" }, "content": { "type": "string" }, "root": root_prop() },
50 - "required": ["path", "content"]
51 - }),
52 - ),
53 - f(
54 - "edit_file",
55 - "Replace an exact, unique occurrence of old_string with new_string in a file.",
56 - json!({
57 - "type": "object",
58 - "properties": {
59 - "path": { "type": "string" },
60 - "old_string": { "type": "string" },
61 - "new_string": { "type": "string" },
62 - "root": root_prop()
63 - },
64 - "required": ["path", "old_string", "new_string"]
65 - }),
66 - ),
67 - f(
68 - "list_directory",
69 - "List entries of a directory (directories end with '/').",
70 - json!({
71 - "type": "object",
72 - "properties": { "path": { "type": "string", "description": "Defaults to the working directory" }, "root": root_prop() }
73 - }),
74 - ),
75 - f(
76 - "run_command",
77 - "Run a shell command (sh -c) in the working directory and return its output.",
78 - json!({
79 - "type": "object",
80 - "properties": { "command": { "type": "string" }, "root": root_prop() },
81 - "required": ["command"]
82 - }),
83 - ),
84 - ]
85 -}
86 -
87 -pub const NAMES: &[&str] = &[
88 - "read_file",
89 - "write_file",
90 - "edit_file",
91 - "list_directory",
92 - "run_command",
93 -];
94 -
95 -pub async fn execute(
96 - ctx: &ToolCtx,
97 - tool_call_id: &str,
98 - name: &str,
99 - args: Value,
100 -) -> Result<ToolOutcome> {
101 - match name {
102 - "read_file" => ctx.read_file(args).await,
103 - "write_file" => ctx.write_file(tool_call_id, args).await,
104 - "edit_file" => ctx.edit_file(tool_call_id, args).await,
105 - "list_directory" => ctx.list_directory(args).await,
106 - "run_command" => ctx.run_command(tool_call_id, args).await,
107 - other => Err(anyhow!("unknown tool `{other}`")),
108 - }
109 -}
110 -
111 -#[derive(Deserialize)]
112 -struct ReadArgs {
113 - path: String,
114 - line: Option<u32>,
115 - limit: Option<u32>,
116 - root: Option<String>,
117 -}
118 -#[derive(Deserialize)]
119 -struct WriteArgs {
120 - path: String,
121 - content: String,
122 - root: Option<String>,
123 -}
124 -#[derive(Deserialize)]
125 -struct EditArgs {
126 - path: String,
127 - old_string: String,
128 - new_string: String,
129 - root: Option<String>,
130 -}
131 -#[derive(Deserialize)]
132 -struct ListArgs {
133 - path: Option<String>,
134 - root: Option<String>,
135 -}
136 -#[derive(Deserialize)]
137 -struct CommandArgs {
138 - command: String,
139 - root: Option<String>,
140 -}
141 -
142 -impl ToolCtx {
143 - pub(super) async fn read_file(&self, args: Value) -> Result<ToolOutcome> {
144 - let a: ReadArgs = serde_json::from_value(args)?;
145 - let path = self.resolve_with(a.root.as_deref(), &a.path)?;
146 - let text = if self.caps.fs.read_text_file {
147 - // ACP v1 fs/read_text_file requires an absolute path.
148 - let mut req = ReadTextFileRequest::new(self.session_id.clone(), absolutize(&path));
149 - if let Some(l) = a.line {
150 - req = req.line(l);
151 - }
152 - if let Some(l) = a.limit {
153 - req = req.limit(l);
154 - }
155 - self.connection
156 - .send_request(req)
157 - .block_task()
158 - .await?
159 - .content
160 - } else {
161 - let full = tokio::fs::read_to_string(&path)
162 - .await
163 - .with_context(|| format!("reading {}", path.display()))?;
164 - let skip = a.line.map_or(0, |l| l.saturating_sub(1) as usize);
165 - let take = a.limit.map_or(usize::MAX, |l| l as usize);
166 - full.lines()
167 - .skip(skip)
168 - .take(take)
169 - .collect::<Vec<_>>()
170 - .join("\n")
171 - };
172 - // Show the user a short summary; give the model the full text.
173 - let mut out = ToolOutcome::ok(format!("{} lines", text.lines().count()));
174 - out.text = truncate(text);
175 - Ok(out)
176 - }
177 -
178 - /// Read the current file contents for diffs/edits; `Ok(None)` if it doesn't exist.
179 - /// A client-fs read failure is an error, distinct from a missing file (ACP
180 - /// `ResourceNotFound`), so callers don't mistake a failed read for a new file.
181 - pub(super) async fn read_existing(&self, path: &Path) -> Result<Option<String>> {
182 - if self.caps.fs.read_text_file {
183 - let req = ReadTextFileRequest::new(self.session_id.clone(), absolutize(path));
184 - match self.connection.send_request(req).block_task().await {
185 - Ok(r) => Ok(Some(r.content)),
186 - Err(e) if e.code == ErrorCode::ResourceNotFound => Ok(None),
187 - Err(e) => Err(anyhow!("client read of {} failed: {e}", path.display())),
188 - }
189 - } else {
190 - match tokio::fs::read_to_string(path).await {
191 - Ok(content) => Ok(Some(content)),
192 - Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
193 - Err(e) => Err(anyhow!("reading {}: {e}", path.display())),
194 - }
195 - }
196 - }
197 -
198 - pub(super) async fn write_text(&self, path: &Path, content: &str) -> Result<()> {
199 - if self.caps.fs.write_text_file {
200 - let req = WriteTextFileRequest::new(self.session_id.clone(), absolutize(path), content);
201 - self.connection.send_request(req).block_task().await?;
202 - } else {
203 - if let Some(parent) = path.parent() {
204 - tokio::fs::create_dir_all(parent).await?;
205 - }
206 - tokio::fs::write(path, content).await?;
207 - }
208 - Ok(())
209 - }
210 -
211 - pub(super) async fn write_file(&self, id: &str, args: Value) -> Result<ToolOutcome> {
212 - let a: WriteArgs = serde_json::from_value(args)?;
213 - let path = self.resolve_with(a.root.as_deref(), &a.path)?;
214 - let old = self.read_existing(&path).await?;
215 - // ACP v1: Diff.path must be an absolute file path.
216 - let diff = Diff::new(absolutize(&path), a.content.clone()).old_text(old);
217 - self.apply_edit(id, "write_file", &path, &a.content, diff)
218 - .await
219 - }
220 -
221 - pub(super) async fn edit_file(&self, id: &str, args: Value) -> Result<ToolOutcome> {
222 - let a: EditArgs = serde_json::from_value(args)?;
223 - let path = self.resolve_with(a.root.as_deref(), &a.path)?;
224 - let old = self
225 - .read_existing(&path)
226 - .await?
227 - .ok_or_else(|| anyhow!("{} does not exist", path.display()))?;
228 - match old.matches(&a.old_string).count() {
229 - 0 => bail!("old_string not found in {}", path.display()),
230 - 1 => {}
231 - n => bail!(
232 - "old_string occurs {n} times in {}; make it unique",
233 - path.display()
234 - ),
235 - }
236 - let new = old.replacen(&a.old_string, &a.new_string, 1);
237 - let diff = Diff::new(absolutize(&path), new.clone()).old_text(old);
238 - self.apply_edit(id, "edit_file", &path, &new, diff).await
239 - }
240 -
241 - pub(super) async fn apply_edit(
242 - &self,
243 - id: &str,
244 - tool: &str,
245 - path: &Path,
246 - content: &str,
247 - diff: Diff,
248 - ) -> Result<ToolOutcome> {
249 - let diff = ToolCallContent::from(diff);
250 - if !self.permit(id, tool, vec![diff.clone()]).await? {
251 - return Ok(ToolOutcome::err("User rejected this change."));
252 - }
253 - self.write_text(path, content).await?;
254 - Ok(ToolOutcome {
255 - text: format!("Wrote {}", path.display()),
256 - content: vec![diff],
257 - failed: false,
258 - locations: Vec::new(),
259 - })
260 - }
261 -
262 - pub(super) async fn list_directory(&self, args: Value) -> Result<ToolOutcome> {
263 - let a: ListArgs = serde_json::from_value(args)?;
264 - let base = self.base_for(a.root.as_deref())?;
265 - let path = a
266 - .path
267 - .map_or_else(|| base.clone(), |p| resolve_in(&base, &p));
268 - let mut entries = Vec::new();
269 - let mut rd = tokio::fs::read_dir(&path)
270 - .await
271 - .with_context(|| format!("listing {}", path.display()))?;
272 - while let Some(e) = rd.next_entry().await? {
273 - let mut name = e.file_name().to_string_lossy().into_owned();
274 - if e.file_type().await.is_ok_and(|t| t.is_dir()) {
275 - name.push('/');
276 - }
277 - entries.push(name);
278 - }
279 - entries.sort();
280 - Ok(ToolOutcome::ok(entries.join("\n")))
281 - }
282 -
283 - pub(super) async fn run_command(&self, id: &str, args: Value) -> Result<ToolOutcome> {
284 - let a: CommandArgs = serde_json::from_value(args)?;
285 - let base = self.base_for(a.root.as_deref())?;
286 - if !self.permit(id, "run_command", vec![]).await? {
287 - return Ok(ToolOutcome::err("User rejected running this command."));
288 - }
289 - if self.caps.terminal {
290 - self.run_in_client_terminal(id, &a.command, &base).await
291 - } else {
292 - self.run_locally(&a.command, &base).await
293 - }
294 - }
295 -
296 - pub(super) async fn run_in_client_terminal(
297 - &self,
298 - id: &str,
299 - command: &str,
300 - cwd: &Path,
301 - ) -> Result<ToolOutcome> {
302 - let sid = self.session_id.clone();
303 - let req = CreateTerminalRequest::new(sid.clone(), "sh")
304 - .args(vec!["-c".into(), command.into()])
305 - // ACP v1: the terminal cwd must be an absolute path.
306 - .cwd(absolutize(cwd))
307 - .output_byte_limit(MAX_OUTPUT_BYTES as u64);
308 - let terminal_id = self
309 - .connection
310 - .send_request(req)
311 - .block_task()
312 - .await?
313 - .terminal_id;
314 - // Embed the live terminal in the tool call so the user can watch it.
315 - let content = vec![ToolCallContent::Terminal(Terminal::new(
316 - terminal_id.clone(),
317 - ))];
318 - self.update(id, ToolCallUpdateFields::new().content(content.clone()))?;
319 -
320 - let wait = self
321 - .connection
322 - .send_request(WaitForTerminalExitRequest::new(
323 - sid.clone(),
324 - terminal_id.clone(),
325 - ))
326 - .block_task();
327 - let exit = tokio::select! {
328 - r = wait => Some(r?.exit_status),
329 - () = tokio::time::sleep(COMMAND_TIMEOUT) => None,
330 - () = self.cancel.cancelled() => None,
331 - };
332 - if exit.is_none() {
333 - let _ = self
334 - .connection
335 - .send_request(KillTerminalRequest::new(sid.clone(), terminal_id.clone()))
336 - .block_task()
337 - .await;
338 - }
339 - let output = self
340 - .connection
341 - .send_request(TerminalOutputRequest::new(sid.clone(), terminal_id.clone()))
342 - .block_task()
343 - .await?;
344 - let _ = self
345 - .connection
346 - .send_request(ReleaseTerminalRequest::new(sid, terminal_id))
347 - .block_task()
348 - .await;
349 -
350 - let status = match &exit {
351 - Some(s) => match (s.exit_code, &s.signal) {
352 - (Some(c), _) => format!("exit code {c}"),
353 - (None, Some(sig)) => format!("killed by {sig}"),
354 - _ => "exited".into(),
355 - },
356 - None if self.cancel.is_cancelled() => "cancelled".into(),
357 - None => format!("timed out after {}s", COMMAND_TIMEOUT.as_secs()),
358 - };
359 - let failed = !matches!(exit.as_ref().and_then(|s| s.exit_code), Some(0));
360 - let trunc = if output.truncated {
361 - "\n[output truncated]"
362 - } else {
363 - ""
364 - };
365 - Ok(ToolOutcome {
366 - text: format!("{}{trunc}\n[{status}]", output.output),
367 - content,
368 - failed,
369 - locations: Vec::new(),
370 - })
371 - }
372 -
373 - pub(super) async fn run_locally(&self, command: &str, cwd: &Path) -> Result<ToolOutcome> {
374 - let child = tokio::process::Command::new("sh")
375 - .arg("-c")
376 - .arg(command)
377 - .current_dir(cwd)
378 - .stdin(std::process::Stdio::null())
379 - .kill_on_drop(true)
380 - .output();
381 - let out = tokio::select! {
382 - r = tokio::time::timeout(COMMAND_TIMEOUT, child) => match r {
383 - Ok(r) => r?,
384 - Err(_) => return Ok(ToolOutcome::err(format!("Command timed out after {}s", COMMAND_TIMEOUT.as_secs()))),
385 - },
386 - () = self.cancel.cancelled() => return Ok(ToolOutcome::err("Command cancelled")),
387 - };
388 - let mut text = String::from_utf8_lossy(&out.stdout).into_owned();
389 - text.push_str(&String::from_utf8_lossy(&out.stderr));
390 - let status = out
391 - .status
392 - .code()
393 - .map_or_else(|| "killed by signal".into(), |c| format!("exit code {c}"));
394 - let text = format!("{}\n[{status}]", truncate(text));
395 - let mut outcome = ToolOutcome::ok(format!("```\n{text}\n```"));
396 - outcome.text = text;
397 - outcome.failed = !out.status.success();
398 - Ok(outcome)
399 - }
400 -}
src/tui.rs deleted
-747
@@ -1,747 +0,0 @@
1 -//! Terminal UI. It is a regular ACP client: it launches this same binary with `--acp` as the
2 -//! agent subprocess and renders the session with ratatui, so it exercises exactly the protocol
3 -//! path an editor would.
4 -
5 -use std::path::{Path, PathBuf};
6 -use std::time::Duration;
7 -
8 -use agent_client_protocol::schema::ProtocolVersion;
9 -use agent_client_protocol::schema::v1::{
10 - CancelNotification, ClientCapabilities, ContentBlock, FileSystemCapabilities,
11 - InitializeRequest, NewSessionRequest, PermissionOptionKind, PromptRequest, ReadTextFileRequest,
12 - ReadTextFileResponse, RequestPermissionOutcome, RequestPermissionRequest,
13 - RequestPermissionResponse, SelectedPermissionOutcome, SessionId, SessionNotification,
14 - SessionUpdate, StopReason, TextContent, ToolCallContent, ToolCallId, ToolCallStatus,
15 - WriteTextFileRequest, WriteTextFileResponse,
16 -};
17 -use agent_client_protocol::{AcpAgent, AcpAgentConfig, Agent, ConnectionTo};
18 -use crossterm::event::{Event, EventStream, KeyCode, KeyEvent, KeyEventKind, KeyModifiers};
19 -use futures::StreamExt;
20 -use ratatui::layout::{Constraint, Layout, Rect};
21 -use ratatui::style::{Color, Modifier, Style, Stylize};
22 -use ratatui::text::{Line, Span};
23 -use ratatui::widgets::{Block, BorderType, Clear, Paragraph, Wrap};
24 -use ratatui::{DefaultTerminal, Frame};
25 -use tokio::sync::{mpsc, oneshot};
26 -
27 -const SPINNER: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
28 -const MAX_TOOL_OUTPUT_LINES: usize = 8;
29 -
30 -enum AppEvent {
31 - Update(SessionUpdate),
32 - Permission(
33 - RequestPermissionRequest,
34 - oneshot::Sender<RequestPermissionOutcome>,
35 - ),
36 - TurnDone(Result<StopReason, String>),
37 -}
38 -
39 -pub async fn run(yolo: bool, extra_roots: Vec<PathBuf>) -> anyhow::Result<()> {
40 - let mut config = AcpAgentConfig::new(std::env::current_exe()?).arg("--acp");
41 - config = config.env("SPLITFIRE_SURFACE", "tui");
42 - if yolo {
43 - config = config.arg("--yolo");
44 - }
45 - let cwd = std::env::current_dir()?;
46 - let llm = crate::llm::LlmConfig::from_env();
47 - let model = llm.model.clone();
48 - let (tx, rx) = mpsc::unbounded_channel();
49 -
50 - agent_client_protocol::Client
51 - .builder()
52 - .on_receive_notification(
53 - {
54 - let tx = tx.clone();
55 - async move |n: SessionNotification, _cx| {
56 - let _ = tx.send(AppEvent::Update(n.update));
57 - Ok(())
58 - }
59 - },
60 - agent_client_protocol::on_receive_notification!(),
61 - )
62 - .on_receive_request(
63 - {
64 - let tx = tx.clone();
65 - async move |req: RequestPermissionRequest, responder, cx| {
66 - // Wait for the user's answer off the dispatch loop so updates keep flowing.
67 - let (answer_tx, answer_rx) = oneshot::channel();
68 - let _ = tx.send(AppEvent::Permission(req, answer_tx));
69 - cx.spawn(async move {
70 - let outcome = answer_rx
71 - .await
72 - .unwrap_or(RequestPermissionOutcome::Cancelled);
73 - responder.respond(RequestPermissionResponse::new(outcome))
74 - })
75 - }
76 - },
77 - agent_client_protocol::on_receive_request!(),
78 - )
79 - // ACP v1 fs/read_text_file: serve file contents from the session cwd,
80 - // honoring the optional 1-based `line` and `limit` params.
81 - .on_receive_request(
82 - async move |req: ReadTextFileRequest, responder, _cx| match tokio::task::spawn_blocking(
83 - {
84 - let path = req.path.clone();
85 - move || read_text_file_blocking(&path, req.line, req.limit)
86 - },
87 - )
88 - .await
89 - {
90 - Ok(Ok(r)) => responder.respond(r),
91 - Ok(Err(e)) => responder.respond_with_error(e),
92 - Err(e) => responder.respond_with_internal_error(e.to_string()),
93 - },
94 - agent_client_protocol::on_receive_request!(),
95 - )
96 - // ACP v1 fs/write_text_file: the client MUST create the file (and parent
97 - // directories) if it doesn't exist.
98 - .on_receive_request(
99 - async move |req: WriteTextFileRequest, responder, _cx| {
100 - match tokio::task::spawn_blocking({
101 - let path = req.path.clone();
102 - let content = req.content.clone();
103 - move || write_text_file_blocking(&path, &content)
104 - })
105 - .await
106 - {
107 - Ok(Ok(r)) => responder.respond(r),
108 - Ok(Err(e)) => responder.respond_with_error(e),
109 - Err(e) => responder.respond_with_internal_error(e.to_string()),
110 - }
111 - },
112 - agent_client_protocol::on_receive_request!(),
113 - )
114 - .connect_with(
115 - AcpAgent::new(config),
116 - |conn: ConnectionTo<Agent>| async move {
117 - conn.send_request(
118 - InitializeRequest::new(ProtocolVersion::V1).client_capabilities(
119 - ClientCapabilities::new().fs(FileSystemCapabilities::new()
120 - .read_text_file(true)
121 - .write_text_file(true)),
122 - ),
123 - )
124 - .block_task()
125 - .await?;
126 - let session = conn
127 - .send_request(
128 - NewSessionRequest::new(cwd.clone())
129 - .additional_directories(extra_roots.clone()),
130 - )
131 - .block_task()
132 - .await?;
133 - let mut app = App::new(conn, session.session_id, tx, model, cwd, extra_roots);
134 - let mut terminal = ratatui::init();
135 - let result = app.run(&mut terminal, rx).await;
136 - ratatui::restore();
137 - result.map_err(|e| {
138 - agent_client_protocol::Error::internal_error().data(format!("{e:#}"))
139 - })
140 - },
141 - )
142 - .await?;
143 - Ok(())
144 -}
145 -
146 -/// `fs/read_text_file` handler body: reads `path`, applying the spec's optional
147 -/// 1-based `line` and `limit` params.
148 -fn read_text_file_blocking(
149 - path: &Path,
150 - line: Option<u32>,
151 - limit: Option<u32>,
152 -) -> agent_client_protocol::Result<ReadTextFileResponse> {
153 - let content = std::fs::read_to_string(path).map_err(|e| {
154 - if e.kind() == std::io::ErrorKind::NotFound {
155 - agent_client_protocol::Error::resource_not_found(Some(path.display().to_string()))
156 - } else {
157 - agent_client_protocol::Error::internal_error().data(e.to_string())
158 - }
159 - })?;
160 - let skip = line.map_or(0, |l| l.saturating_sub(1) as usize);
161 - let take = limit.map_or(usize::MAX, |l| l as usize);
162 - let content = content
163 - .lines()
164 - .skip(skip)
165 - .take(take)
166 - .collect::<Vec<_>>()
167 - .join("\n");
168 - Ok(ReadTextFileResponse::new(content))
169 -}
170 -
171 -/// `fs/write_text_file` handler body: the client MUST create the file if it
172 -/// doesn't exist (spec: https://agentclientprotocol.com/protocol/v1/file-system).
173 -fn write_text_file_blocking(
174 - path: &Path,
175 - content: &str,
176 -) -> agent_client_protocol::Result<WriteTextFileResponse> {
177 - let write = || {
178 - if let Some(parent) = path.parent() {
179 - std::fs::create_dir_all(parent)?;
180 - }
181 - std::fs::write(path, content)
182 - };
183 - write().map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))?;
184 - Ok(WriteTextFileResponse::new())
185 -}
186 -
187 -enum Entry {
188 - User(String),
189 - Agent(String),
190 - Thought(String),
191 - Tool {
192 - id: ToolCallId,
193 - title: String,
194 - status: ToolCallStatus,
195 - output: String,
196 - },
197 - Info(String),
198 - Error(String),
199 -}
200 -
201 -struct App {
202 - conn: ConnectionTo<Agent>,
203 - session_id: SessionId,
204 - tx: mpsc::UnboundedSender<AppEvent>,
205 - model: String,
206 - /// Workspace roots, primary working directory first, then any additional roots.
207 - roots: Vec<PathBuf>,
208 - entries: Vec<Entry>,
209 - input: Vec<char>,
210 - cursor: usize,
211 - busy: bool,
212 - permission: Option<(
213 - RequestPermissionRequest,
214 - oneshot::Sender<RequestPermissionOutcome>,
215 - )>,
216 - /// Lines scrolled up from the bottom of the transcript; 0 follows new output.
217 - scroll_up: usize,
218 - tick: usize,
219 - quit: bool,
220 -}
221 -
222 -impl App {
223 - fn new(
224 - conn: ConnectionTo<Agent>,
225 - session_id: SessionId,
226 - tx: mpsc::UnboundedSender<AppEvent>,
227 - model: String,
228 - cwd: PathBuf,
229 - extra_roots: Vec<PathBuf>,
230 - ) -> Self {
231 - let roots = std::iter::once(cwd.clone())
232 - .chain(extra_roots.into_iter().filter(|r| *r != cwd))
233 - .collect();
234 - Self {
235 - conn,
236 - session_id,
237 - tx,
238 - model,
239 - roots,
240 - entries: vec![Entry::Info(
241 - "Ask for a change, or type /quit to exit.".into(),
242 - )],
243 - input: Vec::new(),
244 - cursor: 0,
245 - busy: false,
246 - permission: None,
247 - scroll_up: 0,
248 - tick: 0,
249 - quit: false,
250 - }
251 - }
252 -
253 - async fn run(
254 - &mut self,
255 - terminal: &mut DefaultTerminal,
256 - mut rx: mpsc::UnboundedReceiver<AppEvent>,
257 - ) -> anyhow::Result<()> {
258 - let mut keys = EventStream::new();
259 - let mut ticker = tokio::time::interval(Duration::from_millis(80));
260 - while !self.quit {
261 - terminal.draw(|f| self.draw(f))?;
262 - tokio::select! {
263 - ev = keys.next() => match ev {
264 - Some(Ok(Event::Key(key))) if key.kind == KeyEventKind::Press => self.on_key(key),
265 - Some(Err(e)) => return Err(e.into()),
266 - None => break,
267 - _ => {}
268 - },
269 - Some(ev) = rx.recv() => self.on_event(ev),
270 - _ = ticker.tick(), if self.busy => self.tick += 1,
271 - }
272 - }
273 - Ok(())
274 - }
275 -
276 - fn on_event(&mut self, ev: AppEvent) {
277 - match ev {
278 - AppEvent::Update(update) => self.on_update(update),
279 - AppEvent::Permission(req, answer) => self.permission = Some((req, answer)),
280 - AppEvent::TurnDone(result) => {
281 - self.busy = false;
282 - match result {
283 - Ok(StopReason::EndTurn) => {}
284 - Ok(StopReason::Cancelled) => {
285 - self.entries.push(Entry::Info("Cancelled.".into()))
286 - }
287 - Ok(other) => self
288 - .entries
289 - .push(Entry::Info(format!("Stopped: {other:?}"))),
290 - Err(e) => self.entries.push(Entry::Error(e)),
291 - }
292 - }
293 - }
294 - }
295 -
296 - fn on_update(&mut self, update: SessionUpdate) {
297 - match update {
298 - SessionUpdate::AgentMessageChunk(chunk) => match self.entries.last_mut() {
299 - Some(Entry::Agent(text)) => text.push_str(&block_text(&chunk.content)),
300 - _ => self.entries.push(Entry::Agent(block_text(&chunk.content))),
301 - },
302 - SessionUpdate::AgentThoughtChunk(chunk) => match self.entries.last_mut() {
303 - Some(Entry::Thought(text)) => text.push_str(&block_text(&chunk.content)),
304 - _ => self
305 - .entries
306 - .push(Entry::Thought(block_text(&chunk.content))),
307 - },
308 - SessionUpdate::ToolCall(call) => self.entries.push(Entry::Tool {
309 - id: call.tool_call_id,
310 - title: call.title,
311 - status: call.status,
312 - output: summarize(&call.content),
313 - }),
314 - SessionUpdate::ToolCallUpdate(update) => {
315 - let entry =
316 - self.entries.iter_mut().rev().find(
317 - |e| matches!(e, Entry::Tool { id, .. } if *id == update.tool_call_id),
318 - );
319 - if let Some(Entry::Tool {
320 - title,
321 - status,
322 - output,
323 - ..
324 - }) = entry
325 - {
326 - let fields = update.fields;
327 - if let Some(s) = fields.status {
328 - *status = s;
329 - }
330 - if let Some(t) = fields.title {
331 - *title = t;
332 - }
333 - if let Some(c) = fields.content {
334 - *output = summarize(&c);
335 - }
336 - }
337 - }
338 - _ => {}
339 - }
340 - }
341 -
342 - fn on_key(&mut self, key: KeyEvent) {
343 - let ctrl = key.modifiers.contains(KeyModifiers::CONTROL);
344 - if self.permission.is_some() {
345 - match key.code {
346 - KeyCode::Char('y') => self.answer_permission(Some(PermissionOptionKind::AllowOnce)),
347 - KeyCode::Char('a') => {
348 - self.answer_permission(Some(PermissionOptionKind::AllowAlways))
349 - }
350 - KeyCode::Char('r') => {
351 - self.answer_permission(Some(PermissionOptionKind::RejectAlways))
352 - }
353 - KeyCode::Char('n') | KeyCode::Esc => {
354 - self.answer_permission(Some(PermissionOptionKind::RejectOnce))
355 - }
356 - KeyCode::Char('c') if ctrl => self.cancel_turn(),
357 - _ => {}
358 - }
359 - return;
360 - }
361 - match key.code {
362 - KeyCode::Char('c') if ctrl => {
363 - if self.busy {
364 - self.cancel_turn();
365 - } else if !self.input.is_empty() {
366 - self.input.clear();
367 - self.cursor = 0;
368 - } else {
369 - self.quit = true;
370 - }
371 - }
372 - KeyCode::Char('d') if ctrl && self.input.is_empty() => self.quit = true,
373 - KeyCode::Char('a') if ctrl => self.cursor = 0,
374 - KeyCode::Char('e') if ctrl => self.cursor = self.input.len(),
375 - KeyCode::Char('u') if ctrl => {
376 - self.input.drain(..self.cursor);
377 - self.cursor = 0;
378 - }
379 - KeyCode::Esc if self.busy => self.cancel_turn(),
380 - KeyCode::Enter if !self.busy => self.submit(),
381 - KeyCode::Char(c) if !ctrl => {
382 - self.input.insert(self.cursor, c);
383 - self.cursor += 1;
384 - }
385 - KeyCode::Backspace if self.cursor > 0 => {
386 - self.cursor -= 1;
387 - self.input.remove(self.cursor);
388 - }
389 - KeyCode::Delete if self.cursor < self.input.len() => {
390 - self.input.remove(self.cursor);
391 - }
392 - KeyCode::Left => self.cursor = self.cursor.saturating_sub(1),
393 - KeyCode::Right => self.cursor = (self.cursor + 1).min(self.input.len()),
394 - KeyCode::Home => self.cursor = 0,
395 - KeyCode::End => self.cursor = self.input.len(),
396 - KeyCode::PageUp => self.scroll_up += 10,
397 - KeyCode::PageDown => self.scroll_up = self.scroll_up.saturating_sub(10),
398 - KeyCode::Up => self.scroll_up += 1,
399 - KeyCode::Down => self.scroll_up = self.scroll_up.saturating_sub(1),
400 - _ => {}
401 - }
402 - }
403 -
404 - fn submit(&mut self) {
405 - let text: String = self.input.iter().collect::<String>().trim().to_string();
406 - if text.is_empty() {
407 - return;
408 - }
409 - self.input.clear();
410 - self.cursor = 0;
411 - if text == "/quit" || text == "/exit" {
412 - self.quit = true;
413 - return;
414 - }
415 - self.entries.push(Entry::User(text.clone()));
416 - self.busy = true;
417 - self.scroll_up = 0;
418 -
419 - let conn = self.conn.clone();
420 - let request = PromptRequest::new(
421 - self.session_id.clone(),
422 - vec![ContentBlock::Text(TextContent::new(text))],
423 - );
424 - let tx = self.tx.clone();
425 - tokio::spawn(async move {
426 - let result = conn.send_request(request).block_task().await;
427 - let _ = tx.send(AppEvent::TurnDone(
428 - result.map(|r| r.stop_reason).map_err(|e| format!("{e}")),
429 - ));
430 - });
431 - }
432 -
433 - fn cancel_turn(&mut self) {
434 - if let Some((_, answer)) = self.permission.take() {
435 - let _ = answer.send(RequestPermissionOutcome::Cancelled);
436 - }
437 - if let Err(e) = self
438 - .conn
439 - .send_notification(CancelNotification::new(self.session_id.clone()))
440 - {
441 - self.entries
442 - .push(Entry::Error(format!("cancel failed: {e}")));
443 - }
444 - }
445 -
446 - fn answer_permission(&mut self, kind: Option<PermissionOptionKind>) {
447 - let Some((req, answer)) = self.permission.take() else {
448 - return;
449 - };
450 - let outcome = req
451 - .options
452 - .iter()
453 - .find(|o| Some(o.kind) == kind)
454 - .map(|o| {
455 - RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(
456 - o.option_id.clone(),
457 - ))
458 - })
459 - .unwrap_or(RequestPermissionOutcome::Cancelled);
460 - let _ = answer.send(outcome);
461 - }
462 -
463 - /// Show paths under a workspace root relative to it; secondary roots get a `name/` prefix.
464 - fn short(&self, text: &str) -> String {
465 - let mut out = text.to_string();
466 - for (i, root) in self.roots.iter().enumerate() {
467 - let mut prefix = root.display().to_string();
468 - prefix.push('/');
469 - let replacement = if i == 0 {
470 - String::new()
471 - } else {
472 - root_name(root)
473 - };
474 - out = out.replace(&prefix, &replacement);
475 - }
476 - out
477 - }
478 -
479 - fn draw(&self, f: &mut Frame) {
480 - let [body, input, status] = Layout::vertical([
481 - Constraint::Min(1),
482 - Constraint::Length(3),
483 - Constraint::Length(1),
484 - ])
485 - .areas(f.area());
486 - self.draw_transcript(f, body);
487 - self.draw_input(f, input);
488 - self.draw_status(f, status);
489 - if let Some((req, _)) = &self.permission {
490 - let title = req.tool_call.fields.title.clone().or_else(|| {
491 - self.entries.iter().rev().find_map(|e| match e {
492 - Entry::Tool { id, title, .. } if *id == req.tool_call.tool_call_id => {
493 - Some(title.clone())
494 - }
495 - _ => None,
496 - })
497 - });
498 - draw_permission(f, &self.short(&title.unwrap_or_else(|| "Run tool".into())));
499 - }
500 - }
501 -
502 - fn draw_transcript(&self, f: &mut Frame, area: Rect) {
503 - let mut lines: Vec<Line> = Vec::new();
504 - for entry in &self.entries {
505 - match entry {
506 - Entry::User(text) => {
507 - for (i, l) in text.lines().enumerate() {
508 - let prefix = if i == 0 { "› " } else { " " };
509 - lines.push(Line::from(vec![prefix.cyan().bold(), l.to_string().bold()]));
510 - }
511 - }
512 - Entry::Agent(text) => {
513 - lines.extend(text.trim().lines().map(|l| Line::from(l.to_string())));
514 - }
515 - Entry::Thought(text) => {
516 - let style = Style::new()
517 - .fg(Color::DarkGray)
518 - .add_modifier(Modifier::ITALIC);
519 - lines.extend(
520 - text.trim()
521 - .lines()
522 - .map(|l| Line::styled(format!(" {l}"), style)),
523 - );
524 - }
525 - Entry::Tool {
526 - title,
527 - status,
528 - output,
529 - ..
530 - } => {
531 - let icon = match status {
532 - ToolCallStatus::Completed => "✓".green(),
533 - ToolCallStatus::Failed => "✗".red(),
534 - _ => SPINNER[self.tick % SPINNER.len()].yellow(),
535 - };
536 - lines.push(Line::from(vec![icon, " ".into(), self.short(title).bold()]));
537 - let output = self.short(output);
538 - let out: Vec<&str> = output
539 - .lines()
540 - .filter(|l| !l.trim().is_empty() && !l.trim_start().starts_with("```"))
541 - .collect();
542 - for l in out.iter().take(MAX_TOOL_OUTPUT_LINES) {
543 - lines.push(Line::from(vec![
544 - " │ ".dark_gray(),
545 - l.to_string().dark_gray(),
546 - ]));
547 - }
548 - if out.len() > MAX_TOOL_OUTPUT_LINES {
549 - let more = out.len() - MAX_TOOL_OUTPUT_LINES;
550 - lines.push(Line::from(format!(" │ … {more} more lines").dark_gray()));
551 - }
552 - }
553 - Entry::Info(text) => lines.push(Line::from(text.clone().dark_gray())),
554 - Entry::Error(text) => lines.push(Line::from(format!("error: {text}").red())),
555 - }
556 - lines.push(Line::default());
557 - }
558 -
559 - let paragraph = Paragraph::new(lines).wrap(Wrap { trim: false });
560 - let total = paragraph.line_count(area.width);
561 - let max_offset = total.saturating_sub(area.height as usize);
562 - let offset = max_offset.saturating_sub(self.scroll_up);
563 - f.render_widget(
564 - paragraph.scroll((offset.min(u16::MAX as usize) as u16, 0)),
565 - area,
566 - );
567 - }
568 -
569 - fn draw_input(&self, f: &mut Frame, area: Rect) {
570 - let border = if self.busy {
571 - Color::DarkGray
572 - } else {
573 - Color::Cyan
574 - };
575 - let block = Block::bordered()
576 - .border_type(BorderType::Rounded)
577 - .border_style(Style::new().fg(border))
578 - .title(" splitfire-agent ");
579 - let inner = block.inner(area);
580 - // Scroll horizontally so the cursor stays visible.
581 - let width = inner.width.saturating_sub(1) as usize;
582 - let start = self.cursor.saturating_sub(width);
583 - let visible: String = self.input[start..].iter().take(width + 1).collect();
584 - let text = if self.input.is_empty() && !self.busy {
585 - Line::from("Ask splitfire-agent to do something…".dark_gray())
586 - } else {
587 - Line::from(visible)
588 - };
589 - f.render_widget(Paragraph::new(text).block(block), area);
590 - if self.permission.is_none() {
591 - f.set_cursor_position((inner.x + (self.cursor - start) as u16, inner.y));
592 - }
593 - }
594 -
595 - fn draw_status(&self, f: &mut Frame, area: Rect) {
596 - let state = if self.permission.is_some() {
597 - "waiting for approval".yellow()
598 - } else if self.busy {
599 - format!("{} working", SPINNER[self.tick % SPINNER.len()]).yellow()
600 - } else {
601 - "ready".green()
602 - };
603 - let hints = if self.busy {
604 - "Esc cancel"
605 - } else {
606 - "Enter send · Ctrl-C quit"
607 - };
608 - let roots = self
609 - .roots
610 - .iter()
611 - .map(|r| r.display().to_string())
612 - .collect::<Vec<_>>()
613 - .join(", ");
614 - let line = Line::from(vec![
615 - " ".into(),
616 - state,
617 - " · ".dark_gray(),
618 - self.model.clone().cyan(),
619 - " · ".dark_gray(),
620 - roots.dark_gray(),
621 - " · ".dark_gray(),
622 - hints.dark_gray(),
623 - " · ↑↓/PgUp/PgDn scroll".dark_gray(),
624 - ]);
625 - f.render_widget(Paragraph::new(line), area);
626 - }
627 -}
628 -
629 -fn draw_permission(f: &mut Frame, title: &str) {
630 - let area = f.area();
631 - let width = (title.chars().count() as u16 + 6).clamp(44, area.width.saturating_sub(4).max(1));
632 - let popup = Rect {
633 - x: area.x + area.width.saturating_sub(width) / 2,
634 - y: area.y + area.height.saturating_sub(7) / 2,
635 - width: width.min(area.width),
636 - height: 7.min(area.height),
637 - };
638 - let block = Block::bordered()
639 - .border_type(BorderType::Rounded)
640 - .border_style(Style::new().fg(Color::Yellow))
641 - .title(" Allow this? ");
642 - let text = vec![
643 - Line::from(title.to_string().bold()),
644 - Line::default(),
645 - Line::from(vec![
646 - Span::from("[y]").green().bold(),
647 - " allow ".into(),
648 - Span::from("[a]").green().bold(),
649 - " always ".into(),
650 - Span::from("[n]").red().bold(),
651 - " reject ".into(),
652 - Span::from("[r]").red().bold(),
653 - " never".into(),
654 - ]),
655 - ];
656 - f.render_widget(Clear, popup);
657 - f.render_widget(
658 - Paragraph::new(text).block(block).wrap(Wrap { trim: false }),
659 - popup,
660 - );
661 -}
662 -
663 -/// Display name for a workspace root: its directory name with a trailing slash.
664 -fn root_name(root: &Path) -> String {
665 - match root.file_name() {
666 - Some(name) => format!("{}/", name.to_string_lossy()),
667 - None => root.display().to_string(),
668 - }
669 -}
670 -
671 -fn block_text(block: &ContentBlock) -> String {
672 - match block {
673 - ContentBlock::Text(t) => t.text.clone(),
674 - _ => String::new(),
675 - }
676 -}
677 -
678 -/// Short, display-friendly summary of tool output.
679 -fn summarize(content: &[ToolCallContent]) -> String {
680 - content
681 - .iter()
682 - .map(|c| match c {
683 - ToolCallContent::Content(c) => block_text(&c.content),
684 - ToolCallContent::Diff(d) => {
685 - let old = d.old_text.as_deref().map_or(0, |t| t.lines().count());
686 - format!(
687 - "{} (-{old} +{} lines)",
688 - d.path.display(),
689 - d.new_text.lines().count()
690 - )
691 - }
692 - _ => String::new(),
693 - })
694 - .filter(|s| !s.is_empty())
695 - .collect::<Vec<_>>()
696 - .join("\n")
697 -}
698 -
699 -#[cfg(test)]
700 -mod fs_handler_tests {
701 - use super::*;
702 -
703 - #[test]
704 - fn read_text_file_slices_by_1_based_line_and_limit() {
705 - let dir = std::env::temp_dir().join("splitfire-agent-tui-fs-test");
706 - std::fs::create_dir_all(&dir).unwrap();
707 - let file = dir.join("lines.txt");
708 - std::fs::write(&file, "one\ntwo\nthree\nfour\n").unwrap();
709 -
710 - let all = read_text_file_blocking(&file, None, None).unwrap();
711 - assert_eq!(all.content, "one\ntwo\nthree\nfour");
712 -
713 - let from_line = read_text_file_blocking(&file, Some(2), None).unwrap();
714 - assert_eq!(from_line.content, "two\nthree\nfour");
715 -
716 - let limited = read_text_file_blocking(&file, Some(2), Some(1)).unwrap();
717 - assert_eq!(limited.content, "two");
718 -
719 - let past_end = read_text_file_blocking(&file, Some(99), None).unwrap();
720 - assert_eq!(past_end.content, "");
721 -
722 - std::fs::remove_dir_all(&dir).ok();
723 - }
724 -
725 - #[test]
726 - fn read_text_file_missing_is_resource_not_found() {
727 - let err = read_text_file_blocking(Path::new("/definitely/does/not/exist.txt"), None, None)
728 - .unwrap_err();
729 - assert_eq!(err.code, agent_client_protocol::ErrorCode::ResourceNotFound);
730 - }
731 -
732 - #[test]
733 - fn write_text_file_creates_missing_file_and_parents() {
734 - let dir = std::env::temp_dir()
735 - .join("splitfire-agent-tui-fs-test-2")
736 - .join("nested");
737 - let file = dir.join("created.txt");
738 - write_text_file_blocking(&file, "hello").unwrap();
739 - assert_eq!(std::fs::read_to_string(&file).unwrap(), "hello");
740 -
741 - // Overwrite existing.
742 - write_text_file_blocking(&file, "again").unwrap();
743 - assert_eq!(std::fs::read_to_string(&file).unwrap(), "again");
744 -
745 - std::fs::remove_dir_all(dir.parent().unwrap()).ok();
746 - }
747 -}
tests/acp_compliance.rs
+33 -1045
@@ -1,293 +1,17 @@
1 -//! ACP v1 compliance and feature tests, driven over raw JSON-RPC against the real binary:
2 -//! version negotiation, auth, session persistence across a killed process, session
3 -//! lifecycle, audio blocks, MCP servers, and the built-in stem separator.
1 +//! SplitFire's ACP tests over raw JSON-RPC against the real binary. The protocol itself (version
2 +//! negotiation, auth, sessions, streaming, approval, MCP) is covered by `ed-acp-testkit`'s
3 +//! conformance suite, run here through [`ed_acp_testkit::conformance_suite!`]. The rest are
4 +//! SplitFire's own: slash commands, audio, the music tools and the built-in stem separator.
5
6 mod common;
7
7 -use std::collections::VecDeque;
8 use std::path::{Path, PathBuf};
9 -use std::process::Stdio;
10 -use std::sync::Arc;
11 -use std::time::Duration;
9
13 -use common::{BASE_TOOL_COUNT, MODEL, MockStats, Script, Step, TEST_KEY, start_mock_llm, temp_dir};
14 -use serde_json::{Value, json};
15 -use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
16 -use tokio::process::{Child, ChildStdin, ChildStdout, Command};
10 +use ed_acp_testkit::{Raw, Step, mock, read_pid, stub_server, temp_dir, tool_text, wait_dead};
11 +use common::{BASE_TOOL_COUNT, target};
12 +use serde_json::json;
13
18 -const BIN: &str = env!("CARGO_BIN_EXE_splitfire-agent");
19 -const STUB: &str = env!("CARGO_BIN_EXE_mcp_stdio_stub");
20 -const TIMEOUT: Duration = Duration::from_secs(30);
21 -
22 -/// What the fake client does when the agent asks for permission.
23 -#[derive(Clone, Copy)]
24 -enum Perm {
25 - Allow,
26 - #[allow(dead_code)]
27 - Reject,
28 - RejectAlways,
29 -}
30 -
31 -struct Raw {
32 - child: Child,
33 - stdin: Option<ChildStdin>,
34 - lines: tokio::io::Lines<BufReader<ChildStdout>>,
35 - next_id: u64,
36 - /// Notifications (`session/update` etc.), in arrival order.
37 - notes: Vec<Value>,
38 - /// Permission requests seen, as the full JSON-RPC `params`.
39 - permissions: Vec<Value>,
40 - perm: Perm,
41 - queued: VecDeque<Value>,
42 -}
43 -
44 -struct Env {
45 - vars: Vec<(String, String)>,
46 -}
47 -
48 -impl Env {
49 - fn new(data_dir: &Path, base_url: &str) -> Self {
50 - let mut e = Env { vars: Vec::new() };
51 - e.set("SPLITFIRE_DATA_DIR", data_dir.display().to_string());
52 - e.set("SPLITFIRE_BASE_URL", base_url);
53 - e.set("SPLITFIRE_MODEL", MODEL);
54 - e.set("ONDE_API_KEY", TEST_KEY);
55 - e.set("SPLITFIRE_MCP_DEMUCS", "off");
56 - // No `demucs` on PATH, wherever the tests run.
57 - e.set("PATH", data_dir.display().to_string());
58 - e
59 - }
60 - fn set(&mut self, k: &str, v: impl Into<String>) -> &mut Self {
61 - self.vars.retain(|(n, _)| n != k);
62 - self.vars.push((k.to_string(), v.into()));
63 - self
64 - }
65 - fn unset(&mut self, k: &str) -> &mut Self {
66 - self.vars.retain(|(n, _)| n != k);
67 - self
68 - }
69 -}
70 -
71 -impl Raw {
72 - fn spawn(env: &Env) -> Raw {
73 - let mut cmd = Command::new(BIN);
74 - cmd.arg("--acp")
75 - .env_clear()
76 - .stdin(Stdio::piped())
77 - .stdout(Stdio::piped())
78 - .stderr(Stdio::inherit())
79 - .kill_on_drop(true);
80 - for (k, v) in &env.vars {
81 - cmd.env(k, v);
82 - }
83 - let mut child = cmd.spawn().expect("spawn agent");
84 - let stdin = child.stdin.take();
85 - let stdout = child.stdout.take().unwrap();
86 - Raw {
87 - child,
88 - stdin,
89 - lines: BufReader::new(stdout).lines(),
90 - next_id: 1,
91 - notes: Vec::new(),
92 - permissions: Vec::new(),
93 - perm: Perm::Allow,
94 - queued: VecDeque::new(),
95 - }
96 - }
97 -
98 - async fn send(&mut self, msg: Value) {
99 - let stdin = self.stdin.as_mut().expect("stdin closed");
100 - stdin
101 - .write_all(format!("{msg}\n").as_bytes())
102 - .await
103 - .unwrap();
104 - stdin.flush().await.unwrap();
105 - }
106 -
107 - async fn notify(&mut self, method: &str, params: Value) {
108 - self.send(json!({"jsonrpc": "2.0", "method": method, "params": params}))
109 - .await;
110 - }
111 -
112 - async fn next_message(&mut self) -> Value {
113 - if let Some(m) = self.queued.pop_front() {
114 - return m;
115 - }
116 - let line = tokio::time::timeout(TIMEOUT, self.lines.next_line())
117 - .await
118 - .expect("timed out waiting for the agent")
119 - .unwrap()
120 - .expect("agent closed stdout");
121 - serde_json::from_str(&line).unwrap_or_else(|e| panic!("non-JSON on stdout ({e}): {line}"))
122 - }
123 -
124 - /// Send a request and wait for its response, serving agent→client requests meanwhile.
125 - async fn call(&mut self, method: &str, params: Value) -> Result<Value, Value> {
126 - let id = self.next_id;
127 - self.next_id += 1;
128 - self.send(json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params}))
129 - .await;
130 - loop {
131 - let msg = self.next_message().await;
132 - if msg.get("method").is_some() {
133 - if msg.get("id").is_some() {
134 - self.serve(&msg).await;
135 - } else {
136 - self.notes.push(msg);
137 - }
138 - continue;
139 - }
140 - if msg["id"] == json!(id) {
141 - return match msg.get("error") {
142 - Some(e) => Err(e.clone()),
143 - None => Ok(msg["result"].clone()),
144 - };
145 - }
146 - }
147 - }
148 -
149 - async fn serve(&mut self, req: &Value) {
150 - let id = req["id"].clone();
151 - let reply = match req["method"].as_str().unwrap() {
152 - "session/request_permission" => {
153 - self.permissions.push(req["params"].clone());
154 - match self.perm {
155 - Perm::Allow => {
156 - json!({"outcome": {"outcome": "selected", "optionId": "allow_once"}})
157 - }
158 - Perm::Reject => {
159 - json!({"outcome": {"outcome": "selected", "optionId": "reject_once"}})
160 - }
161 - Perm::RejectAlways => {
162 - json!({"outcome": {"outcome": "selected", "optionId": "reject_always"}})
163 - }
164 - }
165 - }
166 - other => {
167 - self.send(json!({"jsonrpc": "2.0", "id": id, "error": {"code": -32601, "message": format!("unsupported {other}")}})).await;
168 - return;
169 - }
170 - };
171 - self.send(json!({"jsonrpc": "2.0", "id": id, "result": reply}))
172 - .await;
173 - }
174 -
175 - /// Read whatever the agent sends right after a response (e.g. the commands update that
176 - /// follows `session/new`) until it has been quiet for a moment.
177 - async fn settle(&mut self) {
178 - while let Ok(line) =
179 - tokio::time::timeout(Duration::from_millis(300), self.lines.next_line()).await
180 - {
181 - let Ok(Some(line)) = line else { return };
182 - let msg: Value = serde_json::from_str(&line).unwrap();
183 - if msg.get("id").is_some() && msg.get("method").is_some() {
184 - self.serve(&msg).await;
185 - } else {
186 - self.notes.push(msg);
187 - }
188 - }
189 - }
190 -
191 - async fn ok(&mut self, method: &str, params: Value) -> Value {
192 - self.call(method, params)
193 - .await
194 - .unwrap_or_else(|e| panic!("{method} failed: {e}"))
195 - }
196 -
197 - async fn initialize(&mut self) -> Value {
198 - self.ok(
199 - "initialize",
200 - json!({"protocolVersion": 1, "clientCapabilities": {}}),
201 - )
202 - .await
203 - }
204 -
205 - async fn new_session(&mut self, cwd: &Path, mcp: Value) -> String {
206 - let r = self
207 - .ok("session/new", json!({"cwd": cwd, "mcpServers": mcp}))
208 - .await;
209 - self.settle().await;
210 - r["sessionId"].as_str().unwrap().to_string()
211 - }
212 -
213 - async fn prompt(&mut self, sid: &str, blocks: Value) -> Value {
214 - self.ok(
215 - "session/prompt",
216 - json!({"sessionId": sid, "prompt": blocks}),
217 - )
218 - .await
219 - }
220 -
221 - async fn text_prompt(&mut self, sid: &str, text: &str) -> Value {
222 - self.prompt(sid, json!([{"type": "text", "text": text}]))
223 - .await
224 - }
225 -
226 - /// Notifications for a session update kind, e.g. `tool_call_update`.
227 - fn updates(&self, kind: &str) -> Vec<Value> {
228 - self.notes
229 - .iter()
230 - .filter(|n| {
231 - n["method"] == "session/update" && n["params"]["update"]["sessionUpdate"] == kind
232 - })
233 - .map(|n| n["params"]["update"].clone())
234 - .collect()
235 - }
236 -
237 - /// Close stdin and wait for the agent to exit on its own.
238 - async fn finish(mut self) {
239 - drop(self.stdin.take());
240 - let _ = tokio::time::timeout(TIMEOUT, self.child.wait())
241 - .await
242 - .expect("agent did not exit");
243 - }
244 -}
245 -
246 -fn error_code(e: &Value) -> i64 {
247 - e["code"].as_i64().unwrap()
248 -}
249 -
250 -async fn mock(steps: Vec<Step>) -> (String, Arc<MockStats>) {
251 - let stats = Arc::new(MockStats::default());
252 - let url = start_mock_llm(Script(steps), stats.clone()).await;
253 - (url, stats)
254 -}
255 -
256 -fn tool_text(update: &Value) -> String {
257 - update["content"]
258 - .as_array()
259 - .into_iter()
260 - .flatten()
261 - .filter_map(|c| c.pointer("/content/text").and_then(Value::as_str))
262 - .collect::<Vec<_>>()
263 - .join("\n")
264 -}
265 -
266 -fn alive(pid: u32) -> bool {
267 - std::process::Command::new("kill")
268 - .args(["-0", &pid.to_string()])
269 - .stderr(Stdio::null())
270 - .status()
271 - .is_ok_and(|s| s.success())
272 -}
273 -
274 -async fn wait_dead(pid: u32) -> bool {
275 - for _ in 0..100 {
276 - if !alive(pid) {
277 - return true;
278 - }
279 - tokio::time::sleep(Duration::from_millis(50)).await;
280 - }
281 - false
282 -}
283 -
284 -fn read_pid(file: &Path) -> u32 {
285 - std::fs::read_to_string(file)
286 - .expect("stub pid file")
287 - .trim()
288 - .parse()
289 - .unwrap()
290 -}
14 +ed_acp_testkit::conformance_suite!(target);
15
16 fn wav_bytes() -> Vec<u8> {
17 let mut cursor = std::io::Cursor::new(Vec::new());
@@ -306,600 +30,35 @@ fn wav_bytes() -> Vec<u8> {
30 cursor.into_inner()
31 }
32
309 -fn stub_server(name: &str, pid_file: Option<&Path>) -> Value {
310 - let env: Vec<Value> = pid_file
311 - .map(|p| vec![json!({"name": "STUB_PID_FILE", "value": p})])
312 - .unwrap_or_default();
313 - json!({"name": name, "command": STUB, "args": [], "env": env})
314 -}
315 -
316 -// ---------------------------------------------------------------------------
317 -// initialize / auth
318 -// ---------------------------------------------------------------------------
319 -
320 -#[tokio::test]
321 -async fn initialize_negotiates_v1_and_advertises_capabilities() {
322 - let data = temp_dir("sf-data");
323 - let (url, _) = mock(vec![Step::Final("x")]).await;
324 - let mut agent = Raw::spawn(&Env::new(&data, &url));
325 - // A client from the future gets v1 back, not an echo of its own version.
326 - let r = agent
327 - .ok(
328 - "initialize",
329 - json!({"protocolVersion": 99, "clientCapabilities": {"auth": {"terminal": true}}}),
330 - )
331 - .await;
332 - assert_eq!(r["protocolVersion"], 1);
333 - let caps = &r["agentCapabilities"];
334 - assert_eq!(caps["loadSession"], true);
335 - assert_eq!(caps["promptCapabilities"]["audio"], true);
336 - assert_eq!(caps["promptCapabilities"]["image"], true);
337 - assert_eq!(caps["promptCapabilities"]["embeddedContext"], true);
338 - let sc = &caps["sessionCapabilities"];
339 - for k in ["list", "resume", "close", "delete", "additionalDirectories"] {
340 - assert!(sc.get(k).is_some(), "sessionCapabilities.{k} missing: {sc}");
341 - }
342 - assert_eq!(r["agentInfo"]["name"], "splitfire-agent");
343 - assert!(
344 - caps["auth"].get("logout").is_some(),
345 - "auth.logout missing: {caps}"
346 - );
347 - // Terminal auth is offered to clients that declare they can run it...
348 - assert_eq!(r["authMethods"][0]["id"], "terminal-setup");
349 - assert_eq!(r["authMethods"][0]["args"][0], "--setup");
350 - agent.finish().await;
351 -
352 - let mut agent = Raw::spawn(&Env::new(&data, &url));
353 - let r = agent
354 - .ok(
355 - "initialize",
356 - json!({"protocolVersion": 1, "clientCapabilities": {}}),
357 - )
358 - .await;
359 - assert_eq!(r["protocolVersion"], 1);
360 - // ...and not to clients that can't.
361 - assert!(
362 - r.get("authMethods")
363 - .and_then(Value::as_array)
364 - .is_none_or(|m| m.is_empty()),
365 - "{r}"
366 - );
367 - agent.finish().await;
368 -}
369 -
370 -#[tokio::test]
371 -async fn auth_required_then_authenticate_reloads_config() {
372 - let data = temp_dir("sf-data");
373 - let (url, _) = mock(vec![Step::Final("hello")]).await;
374 - let mut env = Env::new(&data, &url);
375 - env.unset("ONDE_API_KEY");
376 - let mut agent = Raw::spawn(&env);
377 - agent.initialize().await;
378 - let cwd = temp_dir("sf-work");
379 -
380 - // No key anywhere: session/new and session/prompt are refused with auth_required (-32000).
381 - let e = agent
382 - .call("session/new", json!({"cwd": cwd, "mcpServers": []}))
383 - .await
384 - .unwrap_err();
385 - assert_eq!(error_code(&e), -32000, "{e}");
386 - let e = agent
387 - .call("session/prompt", json!({"sessionId": "x", "prompt": []}))
388 - .await
389 - .unwrap_err();
390 - assert_eq!(error_code(&e), -32000, "{e}");
391 - // authenticate fails while there is still no key, and rejects unknown methods.
392 - let e = agent
393 - .call("authenticate", json!({"methodId": "terminal-setup"}))
394 - .await
395 - .unwrap_err();
396 - assert_eq!(error_code(&e), -32000, "{e}");
397 - let e = agent
398 - .call("authenticate", json!({"methodId": "nope"}))
399 - .await
400 - .unwrap_err();
401 - assert_eq!(error_code(&e), -32602, "{e}");
402 -
403 - // load/resume also answer auth_required, so an editor reopening a thread shows sign-in.
404 - let e = agent
405 - .call(
406 - "session/load",
407 - json!({"sessionId": "x", "cwd": cwd, "mcpServers": []}),
408 - )
409 - .await
410 - .unwrap_err();
411 - assert_eq!(error_code(&e), -32000, "{e}");
412 -
413 - // `--setup` stored a key in the config file; authenticate picks it up without a restart.
414 - std::fs::write(data.join("env"), format!("ONDE_API_KEY={TEST_KEY}\n")).unwrap();
415 - agent
416 - .ok("authenticate", json!({"methodId": "terminal-setup"}))
417 - .await;
418 - let sid = agent.new_session(&cwd, json!([])).await;
419 - let r = agent.text_prompt(&sid, "hi").await;
420 - assert_eq!(r["stopReason"], "end_turn");
421 -
422 - // logout removes the stored key; the agent is signed out without a restart.
423 - agent.ok("logout", json!({})).await;
424 - assert!(!data.join("env").exists());
425 - let e = agent
426 - .call("session/new", json!({"cwd": cwd, "mcpServers": []}))
427 - .await
428 - .unwrap_err();
429 - assert_eq!(error_code(&e), -32000, "{e}");
430 -
431 - // Zed's flow: terminal auth runs `--setup` in another process, then retries session/new
432 - // on this one without calling authenticate. The stored key must be picked up anyway.
433 - std::fs::write(data.join("env"), format!("ONDE_API_KEY={TEST_KEY}\n")).unwrap();
434 - agent.new_session(&cwd, json!([])).await;
435 - agent.finish().await;
436 -}
437 -
438 -// ---------------------------------------------------------------------------
439 -// sessions: persistence, replay, lifecycle
440 -// ---------------------------------------------------------------------------
441 -
442 -#[tokio::test]
443 -async fn sessions_survive_a_killed_process_and_replay_tool_output() {
444 - let data = temp_dir("sf-data");
445 - let work = temp_dir("sf-work");
446 - std::fs::write(work.join("notes.txt"), "chorus in Dm").unwrap();
447 - let (url, _) = mock(vec![
448 - Step::ToolCall {
449 - name: "read_file",
450 - arguments: json!({"path": "notes.txt"}),
451 - },
452 - Step::ToolCall {
453 - name: "read_file",
454 - arguments: json!({"path": "missing.txt"}),
455 - },
456 - Step::Final("The chorus is in Dm."),
457 - ])
458 - .await;
459 - let env = Env::new(&data, &url);
460 -
461 - let mut agent = Raw::spawn(&env);
462 - agent.initialize().await;
463 - let sid = agent.new_session(&work, json!([])).await;
464 - let r = agent.text_prompt(&sid, "What key is the chorus in?").await;
465 - assert_eq!(r["stopReason"], "end_turn");
466 - // Title is pushed as a session_info_update.
467 - let infos = agent.updates("session_info_update");
468 - assert!(
469 - infos
470 - .iter()
471 - .any(|u| u["title"] == "What key is the chorus in?"),
472 - "no title update: {infos:?}"
473 - );
474 - // The process dies without a graceful close.
475 - agent.child.kill().await.unwrap();
476 - drop(agent);
477 -
478 - let mut agent = Raw::spawn(&env);
479 - agent.initialize().await;
480 - let list = agent.ok("session/list", json!({})).await;
481 - let sessions = list["sessions"].as_array().unwrap();
482 - assert_eq!(sessions.len(), 1, "{list}");
483 - assert_eq!(sessions[0]["sessionId"], sid.as_str());
484 - assert_eq!(sessions[0]["title"], "What key is the chorus in?");
485 - assert!(sessions[0]["updatedAt"].as_str().unwrap().ends_with('Z'));
486 - // cwd filter.
487 - let other = agent.ok("session/list", json!({"cwd": "/nowhere"})).await;
488 - assert!(other["sessions"].as_array().unwrap().is_empty());
489 -
490 - agent.notes.clear();
491 - agent
492 - .ok(
493 - "session/load",
494 - json!({"sessionId": sid, "cwd": work, "mcpServers": []}),
495 - )
496 - .await;
497 - agent.settle().await;
498 - let user = agent.updates("user_message_chunk");
499 - assert_eq!(user[0]["content"]["text"], "What key is the chorus in?");
500 - let said: String = agent
501 - .updates("agent_message_chunk")
502 - .iter()
503 - .filter_map(|u| u["content"]["text"].as_str())
504 - .collect();
505 - assert!(said.contains("The chorus is in Dm."), "{said}");
506 - let calls = agent.updates("tool_call");
507 - assert_eq!(calls.len(), 2, "{calls:?}");
508 - assert_eq!(calls[0]["status"], "completed");
509 - assert!(
510 - tool_text(&calls[0]).contains("chorus in Dm"),
511 - "{}",
512 - calls[0]
513 - );
514 - assert_eq!(
515 - calls[1]["status"], "failed",
516 - "failed call must replay as failed: {}",
517 - calls[1]
518 - );
519 - assert!(
520 - calls[0]["locations"][0]["path"]
521 - .as_str()
522 - .unwrap()
523 - .ends_with("notes.txt")
524 - );
525 - // Commands are advertised after the load response.
526 - let cmds = agent.updates("available_commands_update");
527 - assert!(!cmds.is_empty());
528 - agent.finish().await;
529 -}
530 -
531 -#[tokio::test]
532 -async fn resume_close_delete_and_unknown_sessions() {
533 - let data = temp_dir("sf-data");
534 - let work = temp_dir("sf-work");
535 - let work2 = temp_dir("sf-work");
536 - let (url, stats) = mock(vec![Step::Final("ok")]).await;
537 - let env = Env::new(&data, &url);
538 - let mut agent = Raw::spawn(&env);
539 - agent.initialize().await;
540 -
541 - let sid = agent.new_session(&work, json!([])).await;
542 - agent.text_prompt(&sid, "first").await;
543 -
544 - // Unknown sessions are invalid_params where nothing can be reopened. load/resume of a
545 - // well-formed id reopen it instead (covered below); a malformed id can't name a session.
546 - for (method, params) in [
547 - (
548 - "session/load",
549 - json!({"sessionId": "../nope", "cwd": work, "mcpServers": []}),
550 - ),
551 - (
552 - "session/resume",
553 - json!({"sessionId": "../nope", "cwd": work, "mcpServers": []}),
554 - ),
555 - (
556 - "session/load",
557 - json!({"sessionId": "relative-cwd", "cwd": "work", "mcpServers": []}),
558 - ),
559 - ("session/new", json!({"cwd": "work", "mcpServers": []})),
560 - ("session/close", json!({"sessionId": "nope"})),
561 - ("session/list", json!({"cursor": "stale"})),
562 - (
563 - "session/prompt",
564 - json!({"sessionId": "nope", "prompt": [{"type": "text", "text": "x"}]}),
565 - ),
566 - (
567 - "session/set_config_option",
568 - json!({"sessionId": "nope", "configId": "model", "value": "other-model"}),
569 - ),
570 - ] {
571 - let e = agent.call(method, params).await.unwrap_err();
572 - assert_eq!(error_code(&e), -32602, "{method}: {e}");
573 - }
574 -
575 - // Config options: model picker lists the endpoint's models; bad ids are rejected.
576 - let created = agent
577 - .ok("session/new", json!({"cwd": work, "mcpServers": []}))
578 - .await;
579 - let opts = created["configOptions"].as_array().unwrap();
580 - assert_eq!(opts[0]["id"], "model");
581 - assert_eq!(opts[0]["category"], "model");
582 - assert_eq!(opts[0]["currentValue"], MODEL);
583 - let values: Vec<&str> = opts[0]["options"]
584 - .as_array()
585 - .unwrap()
586 - .iter()
587 - .map(|o| o["value"].as_str().unwrap())
588 - .collect();
589 - assert!(
590 - values.contains(&MODEL) && values.contains(&"other-model"),
591 - "{values:?}"
592 - );
593 - let sid2 = created["sessionId"].as_str().unwrap().to_string();
594 - let e = agent
595 - .call(
596 - "session/set_config_option",
597 - json!({"sessionId": sid2, "configId": "model", "value": "bogus"}),
598 - )
599 - .await
600 - .unwrap_err();
601 - assert_eq!(error_code(&e), -32602, "{e}");
602 - let e = agent
603 - .call(
604 - "session/set_config_option",
605 - json!({"sessionId": sid2, "configId": "temperature", "value": "x"}),
606 - )
607 - .await
608 - .unwrap_err();
609 - assert_eq!(error_code(&e), -32602, "{e}");
610 - let r = agent
611 - .ok(
612 - "session/set_config_option",
613 - json!({"sessionId": sid2, "configId": "model", "value": "other-model"}),
614 - )
615 - .await;
616 - assert_eq!(r["configOptions"][0]["currentValue"], "other-model");
617 - agent.text_prompt(&sid2, "second").await;
618 - let last = stats.payloads.lock().unwrap().last().cloned().unwrap();
619 - assert_eq!(last["model"], "other-model");
620 -
621 - // Resume rebinds the cwd without replaying anything.
622 - agent.notes.clear();
623 - agent
624 - .ok(
625 - "session/resume",
626 - json!({"sessionId": sid, "cwd": work2, "mcpServers": []}),
627 - )
628 - .await;
629 - assert!(
630 - agent.updates("user_message_chunk").is_empty(),
631 - "resume must not replay"
632 - );
633 - let listed = agent.ok("session/list", json!({"cwd": work2})).await;
634 - assert_eq!(listed["sessions"].as_array().unwrap().len(), 1);
635 -
636 - // Close keeps the stored conversation; delete removes it.
637 - agent.ok("session/close", json!({"sessionId": sid})).await;
638 - let e = agent.call("session/close", json!({"sessionId": sid})).await;
639 - assert!(
640 - e.is_ok(),
641 - "close of a stored-only session is a no-op, not an error"
642 - );
643 - let all = agent.ok("session/list", json!({})).await;
644 - assert_eq!(all["sessions"].as_array().unwrap().len(), 2);
645 - agent.ok("session/delete", json!({"sessionId": sid})).await;
646 - let all = agent.ok("session/list", json!({})).await;
647 - assert_eq!(all["sessions"].as_array().unwrap().len(), 1);
648 - assert!(!data.join("sessions").join(format!("{sid}.json")).exists());
649 - // Delete is idempotent: a second delete, or one for a session that never existed, succeeds.
650 - agent.ok("session/delete", json!({"sessionId": sid})).await;
651 - agent
652 - .ok("session/delete", json!({"sessionId": "nope"}))
653 - .await;
654 -
655 - // A well-formed id with no saved history (a thread from a process that never saved it)
656 - // reopens empty under the same id, with a notice, instead of failing to launch.
657 - agent.notes.clear();
658 - let r = agent
659 - .ok(
660 - "session/load",
661 - json!({"sessionId": "lost-thread", "cwd": work, "mcpServers": []}),
662 - )
663 - .await;
664 - assert_eq!(r["configOptions"][0]["id"], "model");
665 - let notice = agent.updates("agent_message_chunk");
666 - assert_eq!(notice.len(), 1, "{notice:?}");
667 - assert!(
668 - notice[0]["content"]["text"]
669 - .as_str()
670 - .unwrap()
671 - .contains("starts fresh"),
672 - "{notice:?}"
673 - );
674 - assert!(notice[0]["messageId"].is_string());
675 - agent.text_prompt("lost-thread", "hello again").await;
676 - agent.finish().await;
677 -}
678 -
33 +/// Audio links never reach the model as text to read; SplitFire points it at `analyze_audio`.
34 #[tokio::test]
680 -async fn streaming_reports_message_ids_usage_and_inlines_file_links() {
35 +async fn audio_links_point_the_model_at_the_analyzer() {
36 let data = temp_dir("sf-data");
37 let work = temp_dir("sf-work");
683 - std::fs::write(
684 - work.join("chart.md"),
685 - "Verse: Am F C G\nChorus: F G Em Am\nBridge: Dm E\n",
686 - )
687 - .unwrap();
38 std::fs::write(work.join("take.wav"), wav_bytes()).unwrap();
689 - let (url, stats) = mock(vec![Step::Final("Looks like A minor.")]).await;
690 - let mut env = Env::new(&data, &url);
691 - env.set("SPLITFIRE_CONTEXT_WINDOW", "1000");
692 - let mut agent = Raw::spawn(&env);
39 + let (url, stats) = mock(vec![Step::Final("ok")]).await;
40 + let mut agent = Raw::spawn(&target().env(&data, &url));
41 agent.initialize().await;
42 let sid = agent.new_session(&work, json!([])).await;
695 - let chart = format!("file://{}?column=1#L2:2", work.join("chart.md").display());
43 let take = format!("file://{}", work.join("take.wav").display());
44 agent
45 .prompt(
46 &sid,
47 json!([
701 - {"type": "text", "text": "What key is the chorus in?"},
702 - {"type": "resource_link", "name": "chart.md", "uri": chart},
48 + {"type": "text", "text": "What key is this take in?"},
49 {"type": "resource_link", "name": "take.wav", "uri": take},
50 ]),
51 )
52 .await;
707 -
708 - // The selection link is inlined as just that line; the audio link points at the analyzer.
53 let payload = stats.payloads.lock().unwrap()[0].clone();
710 - assert!(payload["stream_options"]["include_usage"] == true);
54 let user = payload["messages"].as_array().unwrap().last().unwrap()["content"]
55 .as_str()
56 .unwrap()
57 .to_string();
715 - assert!(user.contains("Chorus: F G Em Am"), "{user}");
716 - assert!(
717 - !user.contains("Verse:") && !user.contains("Bridge:"),
718 - "{user}"
719 - );
58 assert!(
59 user.contains("Referenced audio file") && user.contains("analyze_audio"),
60 "{user}"
61 );
724 -
725 - // Every chunk of one assistant message carries the same message id.
726 - let chunks = agent.updates("agent_message_chunk");
727 - assert!(!chunks.is_empty());
728 - assert!(
729 - chunks.iter().all(|c| c["messageId"].is_string()),
730 - "{chunks:?}"
731 - );
732 - // Usage from the endpoint, sized by SPLITFIRE_CONTEXT_WINDOW.
733 - let usage = agent.updates("usage_update");
734 - assert_eq!(usage.len(), 1, "{usage:?}");
735 - assert_eq!(
736 - (usage[0]["used"].as_u64(), usage[0]["size"].as_u64()),
737 - (Some(150), Some(1000))
738 - );
739 - agent.finish().await;
740 -}
741 -
742 -#[tokio::test]
743 -async fn permissions_reject_always_raw_output_and_auto_approve() {
744 - let data = temp_dir("sf-data");
745 - let work = temp_dir("sf-work");
746 - let write = |path: &'static str| Step::ToolCall {
747 - name: "write_file",
748 - arguments: json!({"path": path, "content": "x"}),
749 - };
750 - let (url, _) = mock(vec![write("a.txt"), write("b.txt"), Step::Final("done")]).await;
751 - let mut agent = Raw::spawn(&Env::new(&data, &url));
752 - // A client that supports boolean config options gets the auto-approve toggle.
753 - let init = agent
754 - .ok(
755 - "initialize",
756 - json!({"protocolVersion": 1, "clientCapabilities": {"session": {"configOptions": {"boolean": {}}}}}),
757 - )
758 - .await;
759 - assert_eq!(init["protocolVersion"], 1);
760 - let created = agent
761 - .ok("session/new", json!({"cwd": work, "mcpServers": []}))
762 - .await;
763 - agent.settle().await;
764 - let sid = created["sessionId"].as_str().unwrap().to_string();
765 - let opts = created["configOptions"].as_array().unwrap();
766 - assert_eq!(opts[1]["id"], "auto_approve");
767 - assert_eq!(opts[1]["type"], "boolean");
768 - assert_eq!(opts[1]["currentValue"], false);
769 -
770 - // "Always reject" is offered, and once chosen the tool is refused without asking again.
771 - agent.perm = Perm::RejectAlways;
772 - agent.text_prompt(&sid, "write two files").await;
773 - assert_eq!(
774 - agent.permissions.len(),
775 - 1,
776 - "asked again after reject_always"
777 - );
778 - let kinds: Vec<&str> = agent.permissions[0]["options"]
779 - .as_array()
780 - .unwrap()
781 - .iter()
782 - .map(|o| o["kind"].as_str().unwrap())
783 - .collect();
784 - assert!(kinds.contains(&"reject_always"), "{kinds:?}");
785 - assert!(!work.join("a.txt").exists() && !work.join("b.txt").exists());
786 - // Tool calls carry the tool name; updates carry the raw output the model saw.
787 - let calls = agent.updates("tool_call");
788 - assert!(calls.iter().all(|c| c["name"] == "write_file"), "{calls:?}");
789 - let done = agent.updates("tool_call_update");
790 - let failed: Vec<&Value> = done.iter().filter(|u| u["status"] == "failed").collect();
791 - assert_eq!(failed.len(), 2, "{done:?}");
792 - assert_eq!(failed[0]["rawOutput"]["failed"], true);
793 - assert!(
794 - failed[0]["rawOutput"]["output"]
795 - .as_str()
796 - .unwrap()
797 - .contains("rejected")
798 - );
799 -
800 - // Turning auto-approve on is broadcast as a config_option_update and skips the prompts.
801 - agent.notes.clear();
802 - let r = agent
803 - .ok(
804 - "session/set_config_option",
805 - json!({"sessionId": sid, "configId": "auto_approve", "type": "boolean", "value": true}),
806 - )
807 - .await;
808 - assert_eq!(r["configOptions"][1]["currentValue"], true);
809 - let broadcast = agent.updates("config_option_update");
810 - assert_eq!(broadcast.len(), 1, "{broadcast:?}");
811 - let e = agent
812 - .call(
813 - "session/set_config_option",
814 - json!({"sessionId": sid, "configId": "auto_approve", "value": "yes"}),
815 - )
816 - .await
817 - .unwrap_err();
818 - assert_eq!(error_code(&e), -32602, "{e}");
819 -
820 - // A fresh session (no reject_always grant) now writes without asking.
821 - let sid2 = agent.new_session(&work, json!([])).await;
822 - agent
823 - .ok(
824 - "session/set_config_option",
825 - json!({"sessionId": sid2, "configId": "auto_approve", "type": "boolean", "value": true}),
826 - )
827 - .await;
828 - agent.permissions.clear();
829 - agent.text_prompt(&sid2, "write two files").await;
830 - assert!(agent.permissions.is_empty(), "{:?}", agent.permissions);
831 - assert!(work.join("a.txt").exists() && work.join("b.txt").exists());
832 - agent.finish().await;
833 -}
834 -
835 -#[tokio::test]
836 -async fn auto_approve_is_hidden_from_clients_without_boolean_options() {
837 - let data = temp_dir("sf-data");
838 - let work = temp_dir("sf-work");
839 - let (url, _) = mock(vec![Step::Final("x")]).await;
840 - let mut agent = Raw::spawn(&Env::new(&data, &url));
841 - agent.initialize().await;
842 - let created = agent
843 - .ok("session/new", json!({"cwd": work, "mcpServers": []}))
844 - .await;
845 - agent.settle().await;
846 - let opts = created["configOptions"].as_array().unwrap();
847 - assert_eq!(opts.len(), 1, "{opts:?}");
848 - let e = agent
849 - .call(
850 - "session/set_config_option",
851 - json!({"sessionId": created["sessionId"], "configId": "auto_approve", "type": "boolean", "value": true}),
852 - )
853 - .await
854 - .unwrap_err();
855 - assert_eq!(error_code(&e), -32602, "{e}");
856 - agent.finish().await;
857 -}
858 -
859 -#[tokio::test]
860 -async fn content_filter_is_a_refusal_and_private_flags_stay_off_the_wire() {
861 - let data = temp_dir("sf-data");
862 - let work = temp_dir("sf-work");
863 - let (url, stats) = mock(vec![
864 - Step::ToolCall {
865 - name: "read_file",
866 - arguments: json!({"path": "missing.txt"}),
867 - },
868 - Step::ContentFilter,
869 - ])
870 - .await;
871 - let mut agent = Raw::spawn(&Env::new(&data, &url));
872 - agent.initialize().await;
873 - let sid = agent.new_session(&work, json!([])).await;
874 - let r = agent.text_prompt(&sid, "read it").await;
875 - assert_eq!(r["stopReason"], "refusal");
876 - let payloads = stats.payloads.lock().unwrap().clone();
877 - let second = payloads.last().unwrap()["messages"]
878 - .as_array()
879 - .unwrap()
880 - .clone();
881 - let tool = second.iter().find(|m| m["role"] == "tool").unwrap();
882 - assert!(
883 - tool["content"].as_str().unwrap().starts_with("Error:"),
884 - "{tool}"
885 - );
886 - assert!(
887 - tool.get("x_failed").is_none(),
888 - "private field leaked to the model: {tool}"
889 - );
890 - // The system prompt is rebuilt each turn and never stored with the history.
891 - assert_eq!(second[0]["role"], "system");
892 - let stored: Value = serde_json::from_slice(
893 - &std::fs::read(data.join("sessions").join(format!("{sid}.json"))).unwrap(),
894 - )
895 - .unwrap();
896 - assert!(
897 - stored["messages"]
898 - .as_array()
899 - .unwrap()
900 - .iter()
901 - .all(|m| m["role"] != "system")
902 - );
62 agent.finish().await;
63 }
64
@@ -908,7 +67,7 @@ async fn slash_commands_are_advertised_and_expanded() {
67 let data = temp_dir("sf-data");
68 let work = temp_dir("sf-work");
69 let (url, stats) = mock(vec![Step::Final("ok")]).await;
911 - let mut agent = Raw::spawn(&Env::new(&data, &url));
70 + let mut agent = Raw::spawn(&target().env(&data, &url));
71 agent.initialize().await;
72 let sid = agent.new_session(&work, json!([])).await;
73 let cmds = agent.updates("available_commands_update");
@@ -960,7 +119,7 @@ async fn audio_blocks_are_saved_to_scratch_and_cleaned_on_close() {
119 let data = temp_dir("sf-data");
120 let work = temp_dir("sf-work");
121 let (url, stats) = mock(vec![Step::Final("heard it")]).await;
963 - let mut agent = Raw::spawn(&Env::new(&data, &url));
122 + let mut agent = Raw::spawn(&target().env(&data, &url));
123 agent.initialize().await;
124 let sid = agent.new_session(&work, json!([])).await;
125 let b64 = base64::engine::general_purpose::STANDARD.encode(wav_bytes());
@@ -1013,7 +172,7 @@ async fn analyze_audio_runs_through_the_agent() {
172 Step::Final("done"),
173 ])
174 .await;
1016 - let mut agent = Raw::spawn(&Env::new(&data, &url));
175 + let mut agent = Raw::spawn(&target().env(&data, &url));
176 agent.initialize().await;
177 let sid = agent.new_session(&work, json!([])).await;
178 agent.text_prompt(&sid, "analyze tone.wav").await;
@@ -1057,7 +216,7 @@ async fn theory_tools_are_ungated_and_exact() {
216 Step::Final("done"),
217 ])
218 .await;
1060 - let mut agent = Raw::spawn(&Env::new(&data, &url));
219 + let mut agent = Raw::spawn(&target().env(&data, &url));
220 agent.initialize().await;
221 let sid = agent.new_session(&work, json!([])).await;
222 agent.text_prompt(&sid, "C# major scale?").await;
@@ -1081,186 +240,6 @@ async fn theory_tools_are_ungated_and_exact() {
240 agent.finish().await;
241 }
242
1084 -// ---------------------------------------------------------------------------
1085 -// MCP
1086 -// ---------------------------------------------------------------------------
1087 -
1088 -#[tokio::test]
1089 -async fn client_supplied_mcp_server_is_offered_gated_and_streams_progress() {
1090 - let data = temp_dir("sf-data");
1091 - let work = temp_dir("sf-work");
1092 - let pid_file = data.join("stub.pid");
1093 - let (url, stats) = mock(vec![
1094 - Step::ToolCall {
1095 - name: "mcp__stub__echo",
1096 - arguments: json!({"text": "hi"}),
1097 - },
1098 - Step::ToolCall {
1099 - name: "mcp__stub__slow_progress",
1100 - arguments: json!({}),
1101 - },
1102 - Step::Final("all done"),
1103 - ])
1104 - .await;
1105 - let mut agent = Raw::spawn(&Env::new(&data, &url));
1106 - agent.initialize().await;
1107 - let sid = agent
1108 - .new_session(&work, json!([stub_server("stub", Some(&pid_file))]))
1109 - .await;
1110 - let r = agent.text_prompt(&sid, "use the stub").await;
1111 - assert_eq!(r["stopReason"], "end_turn");
1112 -
1113 - // Offered: the namespaced tools ride along with the built-ins.
1114 - let payload = stats.payloads.lock().unwrap()[0].clone();
1115 - let names: Vec<String> = payload["tools"]
1116 - .as_array()
1117 - .unwrap()
1118 - .iter()
1119 - .map(|t| t["function"]["name"].as_str().unwrap().to_string())
1120 - .collect();
1121 - assert!(names.contains(&"mcp__stub__echo".to_string()), "{names:?}");
1122 - assert_eq!(names.len(), BASE_TOOL_COUNT + 5, "{names:?}");
1123 -
1124 - // Gated: only the tool without readOnlyHint asked for permission.
1125 - assert_eq!(agent.permissions.len(), 1, "{:?}", agent.permissions);
1126 - let perm = &agent.permissions[0];
1127 - assert_eq!(perm["toolCall"]["toolCallId"], "call_2");
1128 -
1129 - // Callable, with the result fed back to the model.
1130 - let payloads = stats.payloads.lock().unwrap().clone();
1131 - let tools: Vec<&Value> = payloads[1]["messages"]
1132 - .as_array()
1133 - .unwrap()
1134 - .iter()
1135 - .filter(|m| m["role"] == "tool")
1136 - .collect();
1137 - assert_eq!(tools[0]["content"], "echo: hi");
1138 -
1139 - // Progress: MCP notifications/progress became tool_call_update content.
1140 - let progress: Vec<String> = agent
1141 - .updates("tool_call_update")
1142 - .iter()
1143 - .filter(|u| u["toolCallId"] == "call_2")
1144 - .map(tool_text)
1145 - .collect();
1146 - assert!(
1147 - progress.iter().any(|t| t == "halfway (50%)"),
1148 - "{progress:?}"
1149 - );
1150 - let done = agent
1151 - .updates("tool_call_update")
1152 - .into_iter()
1153 - .rfind(|u| u["toolCallId"] == "call_2")
1154 - .unwrap();
1155 - assert_eq!(done["status"], "completed");
1156 -
1157 - // Children are killed when the session closes.
1158 - let pid = read_pid(&pid_file);
1159 - assert!(alive(pid));
1160 - agent.ok("session/close", json!({"sessionId": sid})).await;
1161 - assert!(
1162 - wait_dead(pid).await,
1163 - "stub server still running after session/close"
1164 - );
1165 - agent.finish().await;
1166 -}
1167 -
1168 -#[tokio::test]
1169 -async fn mcp_children_die_when_the_agent_exits() {
1170 - let data = temp_dir("sf-data");
1171 - let work = temp_dir("sf-work");
1172 - let pid_file = data.join("stub.pid");
1173 - let (url, _) = mock(vec![Step::Final("x")]).await;
1174 - let mut agent = Raw::spawn(&Env::new(&data, &url));
1175 - agent.initialize().await;
1176 - agent
1177 - .new_session(&work, json!([stub_server("stub", Some(&pid_file))]))
1178 - .await;
1179 - let pid = read_pid(&pid_file);
1180 - assert!(alive(pid));
1181 - agent.finish().await;
1182 - assert!(wait_dead(pid).await, "stub server outlived the agent");
1183 -}
1184 -
1185 -#[tokio::test]
1186 -async fn broken_mcp_server_is_not_fatal() {
1187 - let data = temp_dir("sf-data");
1188 - let work = temp_dir("sf-work");
1189 - let (url, stats) = mock(vec![Step::Final("fine")]).await;
1190 - let mut agent = Raw::spawn(&Env::new(&data, &url));
1191 - agent.initialize().await;
1192 - let sid = agent
1193 - .new_session(
1194 - &work,
1195 - json!([
1196 - {"name": "ghost", "command": "/definitely/not/here", "args": [], "env": []},
1197 - {"name": "http", "type": "http", "url": "http://127.0.0.1:1/mcp", "headers": []}
1198 - ]),
1199 - )
1200 - .await;
1201 - let r = agent.text_prompt(&sid, "hello").await;
1202 - assert_eq!(r["stopReason"], "end_turn");
1203 - let payload = stats.payloads.lock().unwrap()[0].clone();
1204 - assert_eq!(payload["tools"].as_array().unwrap().len(), BASE_TOOL_COUNT);
1205 - agent.finish().await;
1206 -}
1207 -
1208 -#[tokio::test]
1209 -async fn cancelling_a_turn_cancels_the_mcp_call() {
1210 - let data = temp_dir("sf-data");
1211 - let work = temp_dir("sf-work");
1212 - let (url, _) = mock(vec![
1213 - Step::ToolCall {
1214 - name: "mcp__stub__wait_for_cancel",
1215 - arguments: json!({}),
1216 - },
1217 - Step::Final("should not get here"),
1218 - ])
1219 - .await;
1220 - let mut agent = Raw::spawn(&Env::new(&data, &url));
1221 - agent.initialize().await;
1222 - let sid = agent
1223 - .new_session(&work, json!([stub_server("stub", None)]))
1224 - .await;
1225 - // Fire the prompt without waiting, then cancel once the call is in flight.
1226 - let id = agent.next_id;
1227 - agent.next_id += 1;
1228 - agent
1229 - .send(
1230 - json!({"jsonrpc": "2.0", "id": id, "method": "session/prompt",
1231 - "params": {"sessionId": sid, "prompt": [{"type": "text", "text": "wait"}]}}),
1232 - )
1233 - .await;
1234 - let mut cancelled = false;
1235 - let stop = loop {
1236 - let msg = agent.next_message().await;
1237 - if msg.get("method").is_some() {
1238 - if msg.get("id").is_some() {
1239 - agent.serve(&msg).await;
1240 - // The permission request means the tool call is about to start.
1241 - tokio::time::sleep(Duration::from_millis(300)).await;
1242 - agent
1243 - .notify("session/cancel", json!({"sessionId": sid}))
1244 - .await;
1245 - cancelled = true;
1246 - } else {
1247 - agent.notes.push(msg);
1248 - }
1249 - } else if msg["id"] == json!(id) {
1250 - break msg["result"]["stopReason"].clone();
1251 - }
1252 - };
1253 - assert!(cancelled, "the call should have asked for permission first");
1254 - assert_eq!(stop, "cancelled");
1255 - let last = agent
1256 - .updates("tool_call_update")
1257 - .into_iter()
1258 - .last()
1259 - .unwrap();
1260 - assert_eq!(last["status"], "failed");
1261 - agent.finish().await;
1262 -}
1263 -
243 // ---------------------------------------------------------------------------
244 // built-in demucs
245 // ---------------------------------------------------------------------------
@@ -1282,9 +261,12 @@ async fn builtin_demucs_appears_only_when_configured() {
261 .await;
262
263 // Configured: SPLITFIRE_DEMUCS_BIN points at a server speaking `demucs --mcp`.
1285 - let mut env = Env::new(&data, &url);
264 + let mut env = target().env(&data, &url);
265 env.set("SPLITFIRE_MCP_DEMUCS", "on");
1287 - env.set("SPLITFIRE_DEMUCS_BIN", STUB);
266 + env.set(
267 + "SPLITFIRE_DEMUCS_BIN",
268 + target().mcp_stub.display().to_string(),
269 + );
270 env.set("STUB_PID_FILE", pid_file.display().to_string());
271 let mut agent = Raw::spawn(&env);
272 agent.initialize().await;
@@ -1350,7 +332,7 @@ async fn builtin_demucs_appears_only_when_configured() {
332 // Unset (no binary, nothing on PATH): no tools, no /stems, and the prompt offers no
333 // separation tool, only how to enable it.
334 for mode in ["off", "on"] {
1353 - let mut env = Env::new(&data, &url);
335 + let mut env = target().env(&data, &url);
336 env.set("SPLITFIRE_MCP_DEMUCS", mode);
337 let (url2, stats2) = mock(vec![Step::Final("x")]).await;
338 env.set("SPLITFIRE_BASE_URL", url2);
@@ -1380,14 +362,20 @@ async fn client_supplied_demucs_wins_over_the_builtin() {
362 let builtin_pid = data.join("builtin.pid");
363 let client_pid = data.join("client.pid");
364 let (url, _) = mock(vec![Step::Final("x")]).await;
1383 - let mut env = Env::new(&data, &url);
365 + let mut env = target().env(&data, &url);
366 env.set("SPLITFIRE_MCP_DEMUCS", "on");
1385 - env.set("SPLITFIRE_DEMUCS_BIN", STUB);
367 + env.set(
368 + "SPLITFIRE_DEMUCS_BIN",
369 + target().mcp_stub.display().to_string(),
370 + );
371 env.set("STUB_PID_FILE", builtin_pid.display().to_string());
372 let mut agent = Raw::spawn(&env);
373 agent.initialize().await;
374 agent
1390 - .new_session(&work, json!([stub_server("demucs", Some(&client_pid))]))
375 + .new_session(
376 + &work,
377 + json!([stub_server(&target().mcp_stub, "demucs", Some(&client_pid))]),
378 + )
379 .await;
380 assert!(client_pid.exists(), "the client's server should run");
381 assert!(
@@ -1442,7 +430,7 @@ async fn real_demucs_end_to_end() {
430 Step::Final("separated"),
431 ])
432 .await;
1445 - let mut env = Env::new(&data, &url);
433 + let mut env = target().env(&data, &url);
434 env.set("SPLITFIRE_MCP_DEMUCS", "on");
435 env.set("SPLITFIRE_DEMUCS_BIN", bin);
436 let mut agent = Raw::spawn(&env);
tests/common/mod.rs
+15 -189
@@ -1,203 +1,29 @@
1 -//! Shared test support: a scripted OpenAI-compatible streaming endpoint and agent launch config.
1 +//! SplitFire's side of the test setup: the conformance [`target`], and launch config for the
2 +//! SDK-driven integration tests. The mock endpoint and the raw client come from `ed-acp-testkit`.
3 #![allow(dead_code)]
4
4 -use std::path::{Path, PathBuf};
5 -use std::sync::atomic::{AtomicUsize, Ordering};
6 -use std::sync::{Arc, Mutex};
5 +use std::path::Path;
6
7 use agent_client_protocol::AcpAgentConfig;
9 -use serde_json::{Value, json};
10 -use tokio::io::{AsyncReadExt, AsyncWriteExt};
11 -use tokio::net::TcpListener;
8
13 -/// Dummy Onde credentials; the mock asserts they arrive as the bearer token.
14 -pub const TEST_KEY: &str = "test-app:test-secret";
15 -pub const MODEL: &str = "scripted-test-model";
9 +#[allow(unused_imports)]
10 +pub use ed_acp_testkit::{MODEL, MockStats, Script, Step, TEST_KEY, Target, start_mock_llm, temp_dir};
11 +
12 /// 5 workspace tools + 6 theory tools + analyze_audio + list_separations.
13 pub const BASE_TOOL_COUNT: usize = 13;
14
19 -#[derive(Debug, Default)]
20 -pub struct MockStats {
21 - pub requests: AtomicUsize,
22 - pub saw_bearer: Mutex<bool>,
23 - /// Every `/chat/completions` request body, in arrival order.
24 - pub payloads: Mutex<Vec<Value>>,
25 -}
26 -
27 -/// What the scripted model does on its next completion.
28 -#[derive(Clone)]
29 -pub enum Step {
30 - /// Emit a tool call, then wait for the tool result.
31 - ToolCall {
32 - name: &'static str,
33 - arguments: Value,
34 - },
35 - /// Emit final assistant text and finish.
36 - Final(&'static str),
37 - /// Stop with `finish_reason: content_filter`.
38 - ContentFilter,
39 -}
40 -
41 -#[derive(Clone)]
42 -pub struct Script(pub Vec<Step>);
43 -
44 -impl Script {
45 - /// Count completed tool results to know how far into the script we are.
46 - fn step_for(&self, messages: &[Value]) -> (usize, Step) {
47 - let done = messages.iter().filter(|m| m["role"] == "tool").count();
48 - (
49 - done,
50 - self.0
51 - .get(done)
52 - .unwrap_or_else(|| self.0.last().expect("empty script"))
53 - .clone(),
54 - )
15 +/// splitfire-agent as the conformance suite sees it: built-in demucs off unless a test opts in.
16 +pub fn target() -> Target {
17 + Target {
18 + bin: env!("CARGO_BIN_EXE_splitfire-agent").into(),
19 + name: "splitfire-agent",
20 + env_prefix: "SPLITFIRE",
21 + mcp_stub: env!("CARGO_BIN_EXE_mcp_stdio_stub").into(),
22 + base_tool_count: BASE_TOOL_COUNT,
23 + env: vec![("SPLITFIRE_MCP_DEMUCS".into(), "off".into())],
24 }
25 }
26
58 -fn sse_chunk(
59 - content: Option<&str>,
60 - tool_call: Option<(&str, &str, &str)>,
61 - finish: Option<&str>,
62 -) -> String {
63 - let mut delta = json!({});
64 - if let Some(c) = content {
65 - delta["content"] = json!(c);
66 - }
67 - if let Some((id, name, args)) = tool_call {
68 - delta["tool_calls"] = json!([{
69 - "index": 0,
70 - "id": id,
71 - "function": { "name": name, "arguments": args },
72 - }]);
73 - }
74 - let chunk = json!({ "choices": [{ "delta": delta, "finish_reason": finish }] });
75 - format!("data: {chunk}\n\n")
76 -}
77 -
78 -async fn read_http_request(socket: &mut tokio::net::TcpStream) -> (String, String) {
79 - let mut buf = Vec::new();
80 - let mut tmp = [0u8; 4096];
81 - let header_end;
82 - let mut content_length = 0usize;
83 - loop {
84 - let n = socket.read(&mut tmp).await.unwrap();
85 - if n == 0 {
86 - panic!("connection closed before headers complete");
87 - }
88 - buf.extend_from_slice(&tmp[..n]);
89 - if let Some(pos) = find_subslice(&buf, b"\r\n\r\n") {
90 - header_end = pos + 4;
91 - let headers = String::from_utf8_lossy(&buf[..pos]).to_string();
92 - for line in headers.lines() {
93 - if let Some(v) = line.to_ascii_lowercase().strip_prefix("content-length:") {
94 - content_length = v.trim().parse().unwrap();
95 - }
96 - }
97 - break;
98 - }
99 - }
100 - while buf.len() < header_end + content_length {
101 - let n = socket.read(&mut tmp).await.unwrap();
102 - if n == 0 {
103 - break;
104 - }
105 - buf.extend_from_slice(&tmp[..n]);
106 - }
107 - let head = String::from_utf8_lossy(&buf[..header_end]).to_string();
108 - let body = String::from_utf8_lossy(&buf[header_end..header_end + content_length]).to_string();
109 - (head, body)
110 -}
111 -
112 -fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option<usize> {
113 - haystack.windows(needle.len()).position(|w| w == needle)
114 -}
115 -
116 -async fn respond(socket: &mut tokio::net::TcpStream, content_type: &str, body: &str) {
117 - let resp = format!(
118 - "HTTP/1.1 200 OK\r\ncontent-type: {content_type}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
119 - body.len()
120 - );
121 - socket.write_all(resp.as_bytes()).await.unwrap();
122 - socket.shutdown().await.unwrap();
123 -}
124 -
125 -/// Spawn the mock server; returns its base URL (e.g. `http://127.0.0.1:PORT/v1`).
126 -pub async fn start_mock_llm(script: Script, stats: Arc<MockStats>) -> String {
127 - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
128 - let port = listener.local_addr().unwrap().port();
129 - tokio::spawn(async move {
130 - loop {
131 - let Ok((mut socket, _)) = listener.accept().await else {
132 - return;
133 - };
134 - let stats = stats.clone();
135 - let script = script.clone();
136 - tokio::spawn(async move {
137 - let (head, body) = read_http_request(&mut socket).await;
138 - let first = head.lines().next().unwrap_or("").to_string();
139 - if first.starts_with("GET /v1/models") {
140 - let list =
141 - json!({"object": "list", "data": [{"id": MODEL}, {"id": "other-model"}]});
142 - return respond(&mut socket, "application/json", &list.to_string()).await;
143 - }
144 - stats.requests.fetch_add(1, Ordering::SeqCst);
145 - if head
146 - .to_ascii_lowercase()
147 - .contains(&format!("authorization: bearer {TEST_KEY}").to_ascii_lowercase())
148 - {
149 - *stats.saw_bearer.lock().unwrap() = true;
150 - }
151 - assert!(
152 - first.starts_with("POST /v1/chat/completions"),
153 - "unexpected request: {first}"
154 - );
155 - let payload: Value = serde_json::from_str(&body).unwrap();
156 - assert_eq!(payload["stream"], true);
157 - stats.payloads.lock().unwrap().push(payload.clone());
158 - let messages = payload["messages"].as_array().unwrap();
159 - let (done, step) = script.step_for(messages);
160 -
161 - let mut data = String::new();
162 - match step {
163 - Step::ToolCall { name, arguments } => {
164 - let id = format!("call_{}", done + 1);
165 - data.push_str(&sse_chunk(
166 - None,
167 - Some((&id, name, &arguments.to_string())),
168 - None,
169 - ));
170 - data.push_str(&sse_chunk(None, None, Some("tool_calls")));
171 - }
172 - Step::Final(text) => {
173 - data.push_str(&sse_chunk(Some(text), None, None));
174 - data.push_str(&sse_chunk(None, None, Some("stop")));
175 - // `stream_options.include_usage`: a final chunk with no choices.
176 - if payload.get("stream_options").is_some() {
177 - let usage = json!({"choices": [], "usage": {
178 - "prompt_tokens": 120, "completion_tokens": 30, "total_tokens": 150
179 - }});
180 - data.push_str(&format!("data: {usage}\n\n"));
181 - }
182 - }
183 - Step::ContentFilter => {
184 - data.push_str(&sse_chunk(None, None, Some("content_filter")));
185 - }
186 - }
187 - data.push_str("data: [DONE]\n\n");
188 - respond(&mut socket, "text/event-stream", &data).await;
189 - });
190 - }
191 - });
192 - format!("http://127.0.0.1:{port}/v1")
193 -}
194 -
195 -pub fn temp_dir(prefix: &str) -> PathBuf {
196 - let dir = std::env::temp_dir().join(format!("{prefix}-{}", uuid::Uuid::new_v4()));
197 - std::fs::create_dir_all(&dir).unwrap();
198 - dir
199 -}
200 -
27 /// Launch config for the real binary: mock endpoint, dummy key, isolated data dir, and no
28 /// built-in demucs unless a test opts in.
29 pub fn agent_config(binary: &str, base_url: &str, data_dir: &Path) -> AcpAgentConfig {