From d1ae121e62ca24a8237e1d254b34caef38819742 Mon Sep 17 00:00:00 2001 From: Yanuar Date: Tue, 13 Oct 2026 10:31:08 +0700 Subject: [PATCH 1/5] fix(acp): compact using provider context usage --- src/acp/service.rs | 599 ++++++++++++++++++++++++++++++++++---- src/session/compaction.rs | 10 +- tests/acp_stdio.rs | 307 ++++++++++++++++++- 3 files changed, 851 insertions(+), 65 deletions(-) diff --git a/src/acp/service.rs b/src/acp/service.rs index 11a35eb..97f3380 100644 --- a/src/acp/service.rs +++ b/src/acp/service.rs @@ -1429,42 +1429,6 @@ impl AcpService { user_message.provider = Some(session.provider.clone()); user_message.model = Some(session.model.clone()); user_message.agent_mode = Some(session.agent.clone()); - let auto_compacted = self - .maybe_auto_compact_session( - &session_id, - &session, - cancellation.clone(), - crate::session::compaction::message_context_tokens(&user_message), - ) - .await?; - if cancellation.is_cancelled() { - return Ok(PromptResponse::new(StopReason::Cancelled)); - } - if auto_compacted { - messages = self - .session_manager - .lock() - .map_err(|_| internal_error())? - .get_session_ref(&session_id) - .map(|stored| stored.messages.clone()) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - } - { - let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; - manager - .add_message_to_session(&session_id, &user_message) - .map_err(|_| internal_error())?; - attachment_guard.commit(); - manager - .set_session_status( - &session_id, - crate::session::types::SessionStatus::Streaming, - None, - ) - .map_err(|_| internal_error())?; - } - messages.push(user_message); - let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); let process_registry = std::sync::Arc::new(crate::tools::ProcessRegistry::new()); let prompt_registry = cancellable( @@ -1504,12 +1468,48 @@ impl AcpService { let Some(system_prompt) = system_prompt else { return Ok(PromptResponse::new(StopReason::Cancelled)); }; - messages.insert(0, crate::session::types::Message::system(system_prompt)); - let base_context_tokens = crate::session::compaction::total_context_tokens(&messages); - let base_cost = messages - .iter() - .filter_map(|message| message.cost) - .sum::(); + let system_message = crate::session::types::Message::system(system_prompt); + let prompt_tokens = crate::session::compaction::message_context_tokens(&system_message); + let auto_compacted = self + .maybe_auto_compact_session( + &session_id, + &session, + &connection, + cancellation.clone(), + prompt_tokens, + crate::session::compaction::message_context_tokens(&user_message), + ) + .await?; + if cancellation.is_cancelled() { + return Ok(PromptResponse::new(StopReason::Cancelled)); + } + if auto_compacted { + messages = self + .session_manager + .lock() + .map_err(|_| internal_error())? + .get_session_ref(&session_id) + .map(|stored| stored.messages.clone()) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + } + { + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + manager + .add_message_to_session(&session_id, &user_message) + .map_err(|_| internal_error())?; + attachment_guard.commit(); + manager + .set_session_status( + &session_id, + crate::session::types::SessionStatus::Streaming, + None, + ) + .map_err(|_| internal_error())?; + } + messages.push(user_message); + let base_context_tokens = acp_context_tokens(&messages, prompt_tokens); + let base_cost = acp_session_cost(&messages); + messages = acp_request_messages(&messages, system_message); send_usage( &connection, &session_id, @@ -1565,12 +1565,20 @@ impl AcpService { &cancellation, receiver, base_context_tokens, + prompt_tokens, base_cost, ) .await?; if !cancellation.is_cancelled() && response.stop_reason != StopReason::Cancelled { let _ = self - .maybe_auto_compact_session(&session_id, &session, cancellation.clone(), 0) + .maybe_auto_compact_session( + &session_id, + &session, + &connection, + cancellation.clone(), + prompt_tokens, + 0, + ) .await?; } Ok(response) @@ -1585,6 +1593,7 @@ impl AcpService { cancellation: &CancellationToken, mut receiver: tokio::sync::mpsc::UnboundedReceiver, base_context_tokens: usize, + base_prompt_tokens: usize, base_cost: f64, ) -> Result { let mut assistant = crate::session::types::Message::incomplete(""); @@ -1693,6 +1702,9 @@ impl AcpService { usage.cache_write, estimate_session_usage_cost(session, &usage), )); + if let Some(part) = assistant.parts.last_mut() { + part.data["acp_prompt_tokens"] = serde_json::json!(base_prompt_tokens); + } if usage.output > 0 { assistant.output_tokens = Some( assistant @@ -1840,17 +1852,7 @@ impl AcpService { match result { Ok(stats) => { - let feedback = format!( - "Context compacted ({})", - crate::session::compaction::format_compaction_stats(stats) - ); - send_text( - &connection, - session_id, - &cuid2::create_id(), - feedback, - false, - )?; + self.notify_compaction(session_id, &session, &connection, stats, 0)?; Ok(PromptResponse::new(StopReason::EndTurn)) } Err(_error) if cancellation.is_cancelled() => { @@ -1959,7 +1961,9 @@ impl AcpService { &self, session_id: &str, session: &AcpSession, + connection: &ConnectionTo, cancellation: CancellationToken, + prompt_tokens: usize, additional_tokens: usize, ) -> Result { let messages = self @@ -1969,8 +1973,8 @@ impl AcpService { .get_session_ref(session_id) .map(|stored| stored.messages.clone()) .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - let used_tokens = crate::session::compaction::total_context_tokens(&messages) - .saturating_add(additional_tokens); + let used_tokens = + acp_context_tokens(&messages, prompt_tokens).saturating_add(additional_tokens); if !crate::session::compaction::should_auto_compact( &session.config.merged_config.compaction, used_tokens, @@ -2018,13 +2022,125 @@ impl AcpService { ) .map_err(|_| internal_error())?; match result { - Ok(_) => Ok(true), + Ok(stats) => { + self.notify_compaction(session_id, session, connection, stats, prompt_tokens)?; + Ok(true) + } Err(error) => { crate::emit_log!("ACP auto-compaction skipped after failure: {:?}", error); Ok(false) } } } + + fn notify_compaction( + &self, + session_id: &str, + session: &AcpSession, + connection: &ConnectionTo, + stats: crate::session::types::CompactionStats, + prompt_tokens: usize, + ) -> Result<(), Error> { + let (used_tokens, cost) = { + let manager = self.session_manager.lock().map_err(|_| internal_error())?; + let stored = manager + .get_session_ref(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + ( + acp_context_tokens(&stored.messages, prompt_tokens), + acp_session_cost(&stored.messages), + ) + }; + send_text( + connection, + session_id, + &cuid2::create_id(), + format!( + "Context compacted ({})", + crate::session::compaction::format_compaction_stats(stats) + ), + false, + )?; + send_usage( + connection, + session_id, + session, + used_tokens, + (cost > 0.0).then_some(cost), + None, + ) + } +} + +fn acp_session_cost(messages: &[crate::session::types::Message]) -> f64 { + messages + .iter() + .map(|message| message.cost.unwrap_or_else(|| message.usage_cost())) + .sum() +} + +fn acp_request_messages( + history: &[crate::session::types::Message], + system_message: crate::session::types::Message, +) -> Vec { + let mut messages = crate::session::compaction::filter_messages_for_context(history); + messages.insert(0, system_message); + messages +} + +fn acp_context_tokens(messages: &[crate::session::types::Message], prompt_tokens: usize) -> usize { + let estimated = + crate::session::compaction::total_context_tokens(messages).saturating_add(prompt_tokens); + // Retained tail usage still describes the old context, as does the summary's billed usage. + let start = messages + .iter() + .rposition(crate::session::compaction::is_compaction_marker) + .map(|index| index + 1) + .unwrap_or(0); + for (index, message) in messages.iter().enumerate().skip(start).rev() { + if message.role != crate::session::types::MessageRole::Assistant { + continue; + } + // Each tool step bills its prompt again; cumulative usage is not a context size. + let provider_context = message + .parts + .iter() + .enumerate() + .filter_map(|(part_index, part)| { + let usage = part.recorded_usage()?; + let input = usage + .input + .saturating_add(usage.cache_read) + .saturating_add(usage.cache_write); + if input == 0 { + return None; + } + let mut tokens = + usize::try_from(input.saturating_add(usage.output)).unwrap_or(usize::MAX); + if let Some(previous_prompt) = part + .data + .get("acp_prompt_tokens") + .and_then(serde_json::Value::as_u64) + { + tokens = tokens + .saturating_sub(usize::try_from(previous_prompt).unwrap_or(usize::MAX)) + .saturating_add(prompt_tokens); + } + Some(tokens.saturating_add( + crate::session::compaction::assistant_parts_context_tokens( + &message.parts[part_index + 1..], + ), + )) + }) + .max() + .filter(|tokens| *tokens > 0); + if let Some(tokens) = provider_context { + let following = + crate::session::compaction::total_context_tokens(&messages[index + 1..]); + return estimated.max(tokens.saturating_add(following)); + } + } + estimated } fn compaction_reasoning(session: &AcpSession) -> Option { @@ -3529,6 +3645,373 @@ mod tests { assert_eq!(live.cumulative.cache_read, 60_000); } + #[test] + fn acp_compaction_counts_provider_context_and_cached_input() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(60_000, 200, 20_000, 5_000, 0.0)); + let messages = vec![Message::user("question"), assistant]; + + assert_eq!(acp_context_tokens(&messages, 1_000), 85_200); + } + + #[test] + fn acp_compaction_does_not_sum_context_across_tool_requests() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(60_000, 200, 20_000, 0, 0.0)); + assistant + .parts + .push(MessagePart::usage(10_000, 100, 40_000, 0, 0.0)); + + assert_eq!( + acp_context_tokens(&[Message::user("question"), assistant], 0), + 80_200 + ); + } + + #[test] + fn acp_compaction_uses_latest_turn_usage_and_counts_pending_history() { + use crate::session::types::{Message, MessagePart}; + let mut old = Message::assistant("old answer"); + old.parts.push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); + let mut latest = Message::assistant("new answer"); + latest + .parts + .push(MessagePart::usage(50_000, 200, 0, 0, 0.0)); + let pending = Message::user("continue"); + let pending_tokens = crate::session::compaction::message_context_tokens(&pending); + let messages = vec![Message::user("question"), old, latest, pending]; + + assert_eq!(acp_context_tokens(&messages, 0), 50_200 + pending_tokens); + } + + #[test] + fn acp_compaction_discards_provider_usage_before_compaction_boundary() { + use crate::session::types::{CompactionStats, Message, MessagePart}; + let mut assistant = Message::assistant("old answer"); + assistant + .parts + .push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); + let mut summary = Message::user(format!( + "{}\nsummary", + crate::session::compaction::SUMMARY_PREFIX + )); + summary + .parts + .push(MessagePart::usage(95_000, 500, 0, 0, 0.0)); + let marker = crate::session::compaction::compaction_marker(CompactionStats { + before_tokens: 90_100, + after_tokens: 100, + before_messages: 2, + after_messages: 1, + }); + let messages = vec![Message::user("question"), assistant, summary, marker]; + + assert_eq!( + acp_context_tokens(&messages, 1_000), + crate::session::compaction::total_context_tokens(&messages) + 1_000 + ); + } + + #[test] + fn acp_compaction_honors_reserved_tokens_and_auto_setting() { + use crate::config::configuration::CompactionConfig; + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(70_000, 0, 10_000, 0, 0.0)); + let messages = vec![Message::user("history ".repeat(2_000)), assistant]; + assert!(crate::session::compaction::select_messages_for_compaction( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + ) + .is_some()); + let config = CompactionConfig::Settings { + auto: true, + prune: false, + reserved: Some(20_000), + }; + let used = acp_context_tokens(&messages, 1_000); + assert!(crate::session::compaction::total_context_tokens(&messages) < 80_000); + assert!(crate::session::compaction::should_auto_compact( + &config, + used, + Some(100_000), + None + )); + assert!(!crate::session::compaction::should_auto_compact( + &config, + used, + Some(100_001), + None + )); + assert!(!crate::session::compaction::should_auto_compact( + &CompactionConfig::Disabled, + used, + Some(100_000), + None + )); + assert!(!crate::session::compaction::should_auto_compact( + &config, used, None, None + )); + let no_auto = CompactionConfig::Settings { + auto: false, + prune: false, + reserved: Some(20_000), + }; + assert!(!crate::session::compaction::should_auto_compact( + &no_auto, + used, + Some(100_000), + None + )); + } + + #[test] + fn acp_compaction_counts_prompt_and_pending_tokens_without_provider_usage() { + use crate::config::configuration::CompactionConfig; + use crate::session::types::Message; + let messages = vec![Message::user("history ".repeat(2_000))]; + let history_tokens = crate::session::compaction::total_context_tokens(&messages); + let config = CompactionConfig::Settings { + auto: true, + prune: false, + reserved: Some(1_000), + }; + let context_window = Some(history_tokens as u32 + 3_000); + assert!(!crate::session::compaction::should_auto_compact( + &config, + acp_context_tokens(&messages, 1_500), + context_window, + None, + )); + assert!(crate::session::compaction::should_auto_compact( + &config, + acp_context_tokens(&messages, 1_500) + 500, + context_window, + None, + )); + } + + #[test] + fn acp_compaction_replaces_previous_prompt_estimate_in_provider_context() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + let mut usage = MessagePart::usage(80_000, 200, 0, 0, 0.0); + usage.data["acp_prompt_tokens"] = serde_json::json!(5_000); + assistant.parts.push(usage); + let messages = vec![Message::user("question"), assistant]; + assert_eq!(acp_context_tokens(&messages, 25_000), 100_200); + assert_eq!(acp_context_tokens(&messages, 1_000), 76_200); + } + + #[test] + fn acp_compaction_keeps_provider_context_when_new_input_usage_is_missing() { + use crate::session::types::{Message, MessagePart}; + for output in [0, 100] { + let mut previous = Message::assistant("previous answer"); + previous + .parts + .push(MessagePart::usage(85_000, 200, 0, 0, 0.0)); + let mut latest = Message::assistant("new answer"); + let mut usage = MessagePart::usage(0, output, 0, 0, 0.0); + usage.data["acp_prompt_tokens"] = serde_json::json!(1_000); + latest.parts.push(usage); + let following_tokens = crate::session::compaction::message_context_tokens(&latest); + assert_eq!( + acp_context_tokens(&[Message::user("question"), previous, latest], 2_000), + 85_200 + following_tokens + ); + } + } + + #[test] + fn acp_compaction_counts_tool_output_after_provider_usage() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(80_000, 200, 0, 0, 0.0)); + record_tool_result( + &mut assistant, + &crate::llm::ToolCallResult { + tool_call_id: "call_1".into(), + role: "tool".into(), + name: "read".into(), + content: serde_json::json!({"output_preview": "late tool output ".repeat(1_000)}) + .to_string(), + }, + ); + let mut following = Message::assistant(""); + following.parts = assistant.parts[2..].to_vec(); + let extra_tokens = crate::session::compaction::message_context_tokens(&following); + assert!(extra_tokens > 1_000); + assert_eq!( + acp_context_tokens(&[Message::user("question"), assistant], 0), + 80_200 + extra_tokens + ); + } + + #[test] + fn acp_compaction_keeps_fresh_system_prompt_in_request() { + use crate::session::types::{CompactionStats, Message}; + let summary = Message::user(format!( + "{}\nsummary", + crate::session::compaction::SUMMARY_PREFIX + )); + let tail = Message::assistant("retained answer"); + let marker = crate::session::compaction::compaction_marker(CompactionStats { + before_tokens: 10_000, + after_tokens: 100, + before_messages: 3, + after_messages: 2, + }); + let pending = Message::user("continue"); + let history = vec![ + Message::user("archived history"), + summary.clone(), + tail.clone(), + marker, + pending.clone(), + ]; + let system = Message::system("fresh system instructions"); + let request = acp_request_messages(&history, system.clone()); + assert_eq!(request, vec![system, summary, tail, pending]); + assert_eq!( + crate::session::compaction::filter_messages_for_context(&request), + request + ); + } + + #[test] + fn acp_compaction_cost_keeps_summary_billing_without_double_counting() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant.cost = Some(0.25); + assistant + .parts + .push(MessagePart::usage(100, 20, 0, 0, 0.25)); + let mut summary = Message::user(format!( + "{}\nsummary", + crate::session::compaction::SUMMARY_PREFIX + )); + summary.parts.push(MessagePart::usage(100, 20, 0, 0, 0.05)); + assert!((acp_session_cost(&[assistant, summary]) - 0.30).abs() < f64::EPSILON); + } + + #[tokio::test(flavor = "current_thread")] + async fn acp_compaction_notifies_client_of_reduced_context_and_summary_cost() { + use crate::session::types::{Message, MessagePart}; + use agent_client_protocol::Agent; + tokio::task::LocalSet::new() + .run_until(async { + let workspace = tempfile::tempdir().unwrap(); + let mut manager = SessionManager::new() + .with_history_for_workspace(workspace.path()) + .unwrap(); + let session_id = manager.create_session(Some("ACP compaction regression".into())); + let mut assistant = Message::assistant("old answer"); + assistant + .parts + .push(MessagePart::usage(80_000, 200, 0, 0, 0.25)); + let messages = vec![Message::user("history ".repeat(2_000)), assistant]; + let selection = crate::session::compaction::select_messages_for_compaction( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + ) + .unwrap(); + let mut session = test_session(); + session.context_window = Some(100_000); + let summary = crate::llm::client::CompactionSummary { + text: "Retained task details".into(), + usage: crate::aisdk::chunk::TokenUsage { + input: 3_000, + output: 20, + cache_read: 0, + cache_write: 0, + }, + }; + let (mut compacted, stats) = compacted_messages( + &messages, + &selection, + &summary, + &session, + crate::session::compaction::total_context_tokens(&messages), + messages.len(), + ) + .unwrap(); + let summary_message = compacted + .iter_mut() + .find(|message| crate::session::compaction::is_compaction_summary(message)) + .unwrap(); + summary_message + .parts + .push(MessagePart::usage(0, 1, 0, 0, 0.05)); + let expected_used = + crate::session::compaction::total_context_tokens(&compacted) + 1_000; + manager + .replace_session_messages(&session_id, compacted) + .unwrap(); + let service = AcpService { + sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), + session_manager: Arc::new(Mutex::new(manager)), + client_capabilities: Arc::new(Mutex::new(Default::default())), + }; + let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); + let client = Client.builder().on_receive_notification( + async move |notification: SessionNotification, _connection| { + let _ = observed_tx.send(notification.update); + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ); + Agent + .builder() + .connect_with(client, async move |connection| { + service + .notify_compaction(&session_id, &session, &connection, stats, 1_000) + .unwrap(); + let feedback = tokio::time::timeout( + std::time::Duration::from_secs(2), + observed_rx.recv(), + ) + .await + .unwrap() + .unwrap(); + assert!(matches!(feedback, SessionUpdate::AgentMessageChunk(_))); + let update = tokio::time::timeout( + std::time::Duration::from_secs(2), + observed_rx.recv(), + ) + .await + .unwrap() + .unwrap(); + let SessionUpdate::UsageUpdate(usage) = update else { + panic!("missing usage update") + }; + assert_eq!(usage.used, expected_used as u64); + assert_eq!(usage.size, 100_000); + assert!((usage.cost.unwrap().amount - 0.30).abs() < f64::EPSILON); + service + .session_manager + .lock() + .unwrap() + .try_delete_session(&session_id) + .unwrap(); + Ok(()) + }) + .await + .unwrap(); + }) + .await; + } + #[test] fn session_info_includes_hierarchy_metadata() { let now = std::time::SystemTime::now(); @@ -4301,7 +4784,7 @@ mod tests { })).unwrap(); } let stream = service.receive_prompt_stream( - &session_id, &session, &connection, &cancellation, receiver, 0, 0.0, + &session_id, &session, &connection, &cancellation, receiver, 0, 0, 0.0, ); tokio::pin!(stream); tokio::time::timeout(Duration::from_secs(2), async { diff --git a/src/session/compaction.rs b/src/session/compaction.rs index af19578..4c84a1e 100644 --- a/src/session/compaction.rs +++ b/src/session/compaction.rs @@ -546,22 +546,20 @@ fn message_parts_context_tokens(message: &Message) -> usize { } match message.role { - MessageRole::Assistant => assistant_parts_context_tokens(message), + MessageRole::Assistant => assistant_parts_context_tokens(&message.parts), MessageRole::Tool => estimate_tokens(&tool_content_for_prompt(&message.content)), _ => estimate_tokens(&message.content), } } -fn assistant_parts_context_tokens(message: &Message) -> usize { - let tool_call_ids = message - .parts +pub(crate) fn assistant_parts_context_tokens(parts: &[MessagePart]) -> usize { + let tool_call_ids = parts .iter() .filter(|part| part.part_type == "tool_call") .filter_map(|part| part.tool_id().map(|id| id.to_string())) .collect::>(); - message - .parts + parts .iter() .map(|part| match part.part_type.as_str() { "text" => part.text_value().map(estimate_tokens).unwrap_or(0), diff --git a/tests/acp_stdio.rs b/tests/acp_stdio.rs index b16988d..4e635d8 100644 --- a/tests/acp_stdio.rs +++ b/tests/acp_stdio.rs @@ -1,8 +1,160 @@ -use std::io::{BufRead, BufReader, Write}; +use std::io::{BufRead, BufReader, Read, Write}; +use std::net::{TcpListener, TcpStream}; use std::process::{Child, ChildStdin, Command, Stdio}; use std::sync::mpsc::{self, Receiver}; use std::time::{Duration, Instant}; +// The provider runs locally, while ACP itself runs in a separate process. No +// process-global environment or external API credentials are needed. +struct MockCompactionProvider { + address: std::net::SocketAddr, + requests: std::sync::Arc>>, + stop: std::sync::Arc, + thread: Option>>, +} + +impl Drop for AcpProcess { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +impl MockCompactionProvider { + fn start() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").expect("mock listener"); + let address = listener.local_addr().unwrap(); + listener.set_nonblocking(true).unwrap(); + let requests = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let recorded = requests.clone(); + let stopping = stop.clone(); + let thread = std::thread::spawn(move || { + let deadline = Instant::now() + Duration::from_secs(60); + let mut regular_turns = 0; + while !stopping.load(std::sync::atomic::Ordering::Relaxed) { + if Instant::now() >= deadline { + return Err("mock provider deadline exceeded".into()); + } + match listener.accept() { + Ok((stream, _)) => { + Self::respond(stream, &recorded, &mut regular_turns) + .map_err(|error| error.to_string())?; + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + std::thread::sleep(Duration::from_millis(5)); + } + Err(error) => return Err(error.to_string()), + } + } + Ok(()) + }); + Self { + address, + requests, + stop, + thread: Some(thread), + } + } + + fn respond( + mut stream: TcpStream, + requests: &std::sync::Mutex>, + regular_turns: &mut usize, + ) -> std::io::Result<()> { + stream.set_read_timeout(Some(Duration::from_secs(5)))?; + stream.set_write_timeout(Some(Duration::from_secs(5)))?; + let mut reader = BufReader::new(stream.try_clone()?); + let mut request_line = String::new(); + reader.read_line(&mut request_line)?; + let mut length = 0; + loop { + let mut header = String::new(); + if reader.read_line(&mut header)? == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "HTTP headers", + )); + } + if header == "\r\n" { + break; + } + if let Some((name, value)) = header.split_once(':') { + if name.eq_ignore_ascii_case("content-length") { + length = value.trim().parse::().expect("HTTP content length"); + } + } + } + let (content_type, body) = if request_line.starts_with("GET /v1/models ") { + ( + "application/json", + serde_json::json!({"object": "list", "data": []}).to_string(), + ) + } else { + assert!( + request_line.starts_with("POST /v1/chat/completions "), + "unexpected request: {request_line}" + ); + assert!(length < 2_000_000, "unexpectedly large mock request"); + let mut body = vec![0; length]; + reader.read_exact(&mut body)?; + let request: serde_json::Value = serde_json::from_slice(&body).expect("provider JSON"); + let compacting = request["messages"] + .to_string() + .contains("CONTEXT CHECKPOINT COMPACTION"); + requests.lock().unwrap().push(request); + let (text, input) = if compacting { + ( + "Checkpoint: the user requested a brief answer. Continue with the next task.", + 300, + ) + } else { + *regular_turns += 1; + ( + "Mock answer.", + if *regular_turns == 1 { 85_000 } else { 500 }, + ) + }; + let chunk = serde_json::json!({ + "id": "mock-completion", "object": "chat.completion.chunk", + "created": 0, "model": "chat", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": null}] + }); + let end = serde_json::json!({ + "id": "mock-completion", "object": "chat.completion.chunk", + "created": 0, "model": "chat", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": input, "completion_tokens": 100, "total_tokens": input + 100} + }); + ( + "text/event-stream", + format!("data: {chunk}\n\ndata: {end}\n\ndata: [DONE]\n\n"), + ) + }; + write!(stream, "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len())?; + stream.flush() + } + + fn finish(mut self) { + self.stop.store(true, std::sync::atomic::Ordering::Relaxed); + self.thread + .take() + .unwrap() + .join() + .expect("mock provider thread") + .expect("mock provider HTTP handling"); + } +} + +impl Drop for MockCompactionProvider { + fn drop(&mut self) { + self.stop.store(true, std::sync::atomic::Ordering::Relaxed); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} + struct AcpProcess { child: Child, stdin: Option, @@ -14,6 +166,13 @@ struct AcpProcess { impl AcpProcess { fn spawn(workspace: &std::path::Path) -> Self { + Self::spawn_with_working_directory(workspace, &std::env::current_dir().unwrap()) + } + + fn spawn_with_working_directory( + workspace: &std::path::Path, + working_directory: &std::path::Path, + ) -> Self { let home = tempfile::tempdir().expect("home"); let config = tempfile::tempdir().expect("config"); let state = tempfile::tempdir().expect("state"); @@ -33,6 +192,7 @@ impl AcpProcess { let mut child = Command::new(env!("CARGO_BIN_EXE_crabcode")) .args(["acp", "--cwd"]) .arg(workspace) + .current_dir(working_directory) .env("HOME", home.path()) .env("XDG_CONFIG_HOME", config.path()) .env("XDG_STATE_HOME", state.path()) @@ -320,6 +480,23 @@ fn session_load_error_keeps_stdio_transport_open() { #[test] fn advertises_and_dispatches_commands_over_stdio() { let workspace = tempfile::tempdir().expect("workspace"); + // The empty models.dev cache avoids external catalog requests, but session + // setup still needs a configured model. These commands never call a provider; + // omitting baseURL also avoids custom-provider model discovery requests. + std::fs::write( + workspace.path().join("crabcode.json"), + serde_json::json!({ + "model": "offline/chat", + "enabled_providers": ["offline"], + "provider": {"offline": { + "npm": "@ai-sdk/openai-compatible", + "options": {"apiKey": "dummy"}, + "models": {"chat": {"limit": {"context": 100000, "output": 1000}}} + }} + }) + .to_string(), + ) + .expect("offline workspace config"); let command_dir = workspace.path().join(".crabcode/commands"); std::fs::create_dir_all(&command_dir).expect("command dir"); std::fs::write( @@ -441,3 +618,131 @@ Wire skill instructions. process.close_and_wait(); } + +#[test] +fn automatically_compacts_provider_usage_over_stdio() { + let workspace = tempfile::tempdir().expect("workspace"); + let provider = MockCompactionProvider::start(); + std::fs::write( + workspace.path().join("crabcode.json"), + serde_json::json!({ + "model": "mock/chat", + "enabled_providers": ["mock"], + "provider": {"mock": { + "npm": "@ai-sdk/openai-compatible", + "options": {"baseURL": format!("http://{}/v1", provider.address), "apiKey": "dummy"}, + "models": {"chat": {"limit": {"context": 100000, "output": 1000}}} + }}, + "compaction": {"auto": true, "reserved": 20000} + }).to_string(), + ).expect("workspace config"); + let system_marker = "ACP_MOCK_WORKSPACE_INSTRUCTIONS_PRESERVED"; + std::fs::write(workspace.path().join("AGENTS.md"), system_marker).unwrap(); + // Request-time provider discovery also reads the process working directory. + let mut process = AcpProcess::spawn_with_working_directory(workspace.path(), workspace.path()); + assert!(initialize(&mut process).get("error").is_none()); + process.send(serde_json::json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", + "params": {"cwd": workspace.path(), "mcpServers": []} + })); + let (response, _) = process.recv_response(2); + let session_id = response["result"]["sessionId"] + .as_str() + .expect("session id"); + // 3,000 estimated tokens clears the 2,000-token compaction minimum but + // cannot reach the 80,000-token threshold without authoritative usage. + let initial_prompt = format!( + "ARCHIVED_LARGE_INITIAL_PROMPT {}Please answer briefly.", + "word ".repeat(2400) + ); + assert!(initial_prompt.len() / 4 >= 2000); + assert!(initial_prompt.len() / 4 < 80000); + process.send(serde_json::json!({ + "jsonrpc": "2.0", "id": 3, "method": "session/prompt", + "params": {"sessionId": session_id, "prompt": [{"type": "text", "text": initial_prompt}]} + })); + let (response, notifications) = process.recv_response(3); + assert_eq!(response["result"]["stopReason"], "end_turn", "{response}"); + let compacted = notifications + .iter() + .position(|message| { + message["params"]["update"]["content"]["text"] + .as_str() + .is_some_and(|text| text.contains("Context compacted")) + }) + .expect("automatic compaction notification without /compact"); + let usage = notifications + .iter() + .filter(|message| message["params"]["update"]["sessionUpdate"] == "usage_update") + .collect::>(); + assert!( + usage + .iter() + .any(|message| message["params"]["update"]["used"] + .as_u64() + .is_some_and(|used| used >= 85000)), + "provider usage must reach ACP: {usage:?}" + ); + let final_usage = &usage.last().expect("final usage")["params"]["update"]; + assert_eq!(final_usage["size"], 100000); + assert!( + final_usage["used"].as_u64().unwrap() < 80000, + "{final_usage}" + ); + assert!(notifications[compacted + 1..] + .iter() + .any(|message| message["params"]["update"]["sessionUpdate"] == "usage_update")); + { + let requests = provider.requests.lock().unwrap(); + assert_eq!( + requests.len(), + 2, + "one regular request then one summary request" + ); + assert!( + requests[0]["messages"].to_string().chars().count() / 4 < 80_000, + "estimated initial transcript alone must not trigger compaction" + ); + assert!(!requests[0]["messages"] + .to_string() + .contains("CONTEXT CHECKPOINT COMPACTION")); + assert!(requests[1]["messages"] + .to_string() + .contains("CONTEXT CHECKPOINT COMPACTION")); + } + process.send(serde_json::json!({ + "jsonrpc": "2.0", "id": 4, "method": "session/prompt", + "params": {"sessionId": session_id, "prompt": [{"type": "text", "text": "Now give another brief answer."}]} + })); + let (response, notifications) = process.recv_response(4); + assert_eq!(response["result"]["stopReason"], "end_turn", "{response}"); + assert!( + !notifications + .iter() + .any(|message| message["params"]["update"]["content"]["text"] + .as_str() + .is_some_and(|text| text.contains("Context compacted"))), + "stale usage must not compact the second turn" + ); + { + let requests = provider.requests.lock().unwrap(); + assert_eq!(requests.len(), 3, "no repeated compaction from stale usage"); + let messages = requests[2]["messages"] + .as_array() + .expect("second-turn messages"); + assert_eq!(messages[0]["role"], "system"); + assert!( + messages[0]["content"].to_string().contains(system_marker), + "fresh system prompt preserves workspace instructions" + ); + let text = requests[2]["messages"].to_string(); + assert!(text.contains("Checkpoint:"), "summary is active context"); + assert!( + !text.contains("ARCHIVED_LARGE_INITIAL_PROMPT"), + "archived transcript must not be resent" + ); + assert!(!text.contains("CONTEXT CHECKPOINT COMPACTION")); + } + process.close_and_wait(); + provider.finish(); +} From 4a4f25a415d14a8e6cdeded0bf482a7abcad7b2c Mon Sep 17 00:00:00 2001 From: Yanuar Date: Tue, 13 Oct 2026 11:44:43 +0700 Subject: [PATCH 2/5] fix(acp): describe configured compaction mode in picker --- src/acp/service.rs | 84 ++++++++++++++++++++++++++++++++++++++++++++-- tests/acp_stdio.rs | 23 +++++++++++++ 2 files changed, 105 insertions(+), 2 deletions(-) diff --git a/src/acp/service.rs b/src/acp/service.rs index 97f3380..e0fa41d 100644 --- a/src/acp/service.rs +++ b/src/acp/service.rs @@ -2262,6 +2262,14 @@ fn tool_permissions(session: &AcpSession) -> crate::tools::ToolPermissions { session.tool_permissions.clone() } +fn mode_description(agent: &crate::agent::definition::AgentDefinition) -> &str { + if agent.name == "compaction" && agent.description.trim().is_empty() { + "Compaction agent. Summarizes conversation context for continued work." + } else { + &agent.description + } +} + fn session_modes(session: &AcpSession) -> SessionModeState { let mut modes = session .config @@ -2271,7 +2279,7 @@ fn session_modes(session: &AcpSession) -> SessionModeState { .into_iter() .map(|agent| { SessionMode::new(agent.name.clone(), agent.name.clone()) - .description(agent.description.clone()) + .description(mode_description(agent)) }) .collect::>(); modes.sort_by(|left, right| left.name.cmp(&right.name)); @@ -2287,7 +2295,7 @@ fn session_config_options(session: &AcpSession) -> Vec { .into_iter() .map(|agent| { SessionConfigSelectOption::new(agent.name.clone(), agent.name.clone()) - .description(agent.description.clone()) + .description(mode_description(agent)) }) .collect::>(); mode_options.sort_by(|left, right| left.name.cmp(&right.name)); @@ -5114,6 +5122,78 @@ mod tests { assert!(workspace_path(Path::new("relative")).is_err()); } + #[test] + fn compaction_mode_descriptions_preserve_configuration_and_visibility() { + let mut session = test_session(); + assert!(!session_modes(&session) + .available_modes + .iter() + .any(|mode| mode.id.to_string() == "compaction")); + + for description in ["", " \n ", "Use this agent to preserve project decisions."] { + let definitions = crate::agent::definition::parse_agent_definitions_from_config( + Some(&serde_json::json!({"compaction": { + "mode": "primary", "description": description + }})), + &mut Vec::new(), + ); + session.config.merged_config.agent_registry = + crate::agent::definition::AgentRegistry::with_definitions(None, definitions); + let expected = if description.trim().is_empty() { + "Compaction agent. Summarizes conversation context for continued work." + } else { + description + }; + let modes = session_modes(&session); + let compaction = modes + .available_modes + .iter() + .find(|mode| mode.id.to_string() == "compaction") + .expect("configured compaction mode"); + assert_eq!(compaction.description.as_deref(), Some(expected)); + let config = serde_json::to_value(session_config_options(&session)).unwrap(); + let compaction = config[0]["options"] + .as_array() + .unwrap() + .iter() + .find(|option| option["value"] == "compaction") + .expect("configured compaction option"); + assert_eq!(compaction["description"], expected); + assert_eq!( + modes + .available_modes + .iter() + .find(|mode| mode.id.to_string() == "build") + .unwrap() + .description + .as_deref(), + Some("The default agent. Executes tools based on configured permissions.") + ); + } + + for settings in [ + serde_json::json!({"mode": "subagent"}), + serde_json::json!({"mode": "primary", "hidden": true}), + ] { + let definitions = crate::agent::definition::parse_agent_definitions_from_config( + Some(&serde_json::json!({"compaction": settings})), + &mut Vec::new(), + ); + session.config.merged_config.agent_registry = + crate::agent::definition::AgentRegistry::with_definitions(None, definitions); + assert!(!session_modes(&session) + .available_modes + .iter() + .any(|mode| mode.id.to_string() == "compaction")); + let config = serde_json::to_value(session_config_options(&session)).unwrap(); + assert!(!config[0]["options"] + .as_array() + .unwrap() + .iter() + .any(|option| option["value"] == "compaction")); + } + } + #[test] fn builds_grouped_model_config_option() { let option = model_config_option( diff --git a/tests/acp_stdio.rs b/tests/acp_stdio.rs index 4e635d8..26a17a2 100644 --- a/tests/acp_stdio.rs +++ b/tests/acp_stdio.rs @@ -488,6 +488,7 @@ fn advertises_and_dispatches_commands_over_stdio() { serde_json::json!({ "model": "offline/chat", "enabled_providers": ["offline"], + "agent": {"compaction": {"mode": "primary"}}, "provider": {"offline": { "npm": "@ai-sdk/openai-compatible", "options": {"apiKey": "dummy"}, @@ -534,6 +535,28 @@ Wire skill instructions. } })); let (new_session, mut notifications) = process.recv_response(2); + let expected_description = + "Compaction agent. Summarizes conversation context for continued work."; + let compaction_mode = new_session["result"]["modes"]["availableModes"] + .as_array() + .expect("available modes") + .iter() + .find(|mode| mode["id"] == "compaction") + .expect("compaction mode"); + assert_eq!(compaction_mode["description"], expected_description); + let mode_config = new_session["result"]["configOptions"] + .as_array() + .expect("config options") + .iter() + .find(|option| option["id"] == "mode") + .expect("mode config"); + let compaction_option = mode_config["options"] + .as_array() + .expect("mode options") + .iter() + .find(|option| option["value"] == "compaction") + .expect("compaction option"); + assert_eq!(compaction_option["description"], expected_description); let session_id = new_session["result"]["sessionId"] .as_str() .expect("session id") From d81fa25829c930a129de84ff7d8ee48d3b777238 Mon Sep 17 00:00:00 2001 From: Yanuar Date: Tue, 13 Oct 2026 11:37:20 +0700 Subject: [PATCH 3/5] refactor(acp): split service into focused modules --- src/acp/service.rs | 5549 ------------------------------ src/acp/service/attachments.rs | 400 +++ src/acp/service/commands.rs | 465 +++ src/acp/service/compaction.rs | 610 ++++ src/acp/service/config.rs | 600 ++++ src/acp/service/mod.rs | 190 + src/acp/service/notifications.rs | 622 ++++ src/acp/service/permissions.rs | 404 +++ src/acp/service/prompt.rs | 367 ++ src/acp/service/questions.rs | 337 ++ src/acp/service/sessions.rs | 411 +++ src/acp/service/stream.rs | 550 +++ src/acp/service/terminals.rs | 371 ++ src/acp/service/test_support.rs | 102 + src/acp/service/usage.rs | 474 +++ 15 files changed, 5903 insertions(+), 5549 deletions(-) delete mode 100644 src/acp/service.rs create mode 100644 src/acp/service/attachments.rs create mode 100644 src/acp/service/commands.rs create mode 100644 src/acp/service/compaction.rs create mode 100644 src/acp/service/config.rs create mode 100644 src/acp/service/mod.rs create mode 100644 src/acp/service/notifications.rs create mode 100644 src/acp/service/permissions.rs create mode 100644 src/acp/service/prompt.rs create mode 100644 src/acp/service/questions.rs create mode 100644 src/acp/service/sessions.rs create mode 100644 src/acp/service/stream.rs create mode 100644 src/acp/service/terminals.rs create mode 100644 src/acp/service/test_support.rs create mode 100644 src/acp/service/usage.rs diff --git a/src/acp/service.rs b/src/acp/service.rs deleted file mode 100644 index e0fa41d..0000000 --- a/src/acp/service.rs +++ /dev/null @@ -1,5549 +0,0 @@ -use crate::config::configuration::LoadedConfig; -use crate::session::manager::SessionManager; -use agent_client_protocol::schema::v1::{ - AvailableCommand, AvailableCommandInput, AvailableCommandsUpdate, ContentBlock, ContentChunk, - Cost as AcpCost, CreateElicitationRequest, CreateTerminalRequest, ElicitationAction, - ElicitationContentValue, ElicitationFormMode, ElicitationSchema, ElicitationSessionScope, - EmbeddedResourceResource, EnumOption, KillTerminalRequest, ListSessionsResponse, - LoadSessionResponse, McpServer, MultiSelectPropertySchema, NewSessionResponse, - PermissionOption, PermissionOptionKind, PromptResponse, ReleaseTerminalRequest, - RequestPermissionOutcome, RequestPermissionRequest, ResumeSessionResponse, SessionConfigOption, - SessionConfigOptionCategory, SessionConfigSelectGroup, SessionConfigSelectOption, SessionInfo, - SessionMode, SessionModeState, SessionNotification, SessionUpdate, - SetSessionConfigOptionResponse, StopReason, StringPropertySchema, Terminal, - TerminalOutputRequest, ToolCall, ToolCallContent, ToolCallLocation, ToolCallStatus, - ToolCallUpdate, ToolCallUpdateFields, ToolKind, UnstructuredCommandInput, UsageUpdate, - WaitForTerminalExitRequest, -}; -use agent_client_protocol::{Client, ConnectionTo, Error}; -use base64::Engine as _; -use std::collections::{HashMap, HashSet}; -use std::path::{Path, PathBuf}; -use std::sync::{Arc, Mutex}; -use tokio::sync::Mutex as AsyncMutex; -use tokio_util::sync::CancellationToken; - -#[derive(Clone)] -pub struct AcpService { - sessions: Arc>, - session_manager: Arc>, - client_capabilities: Arc>, -} - -// Attachments detach immediately on close, but a turn retains ownership until all -// transcript/status finalization finishes. Both are protected by the same lock. -#[derive(Default)] -struct AcpSessions { - attachments: HashMap, - active_prompts: HashSet, -} - -impl AcpSessions { - fn ensure_attachable(&self, session_id: &str) -> Result<(), Error> { - if self.active_prompts.contains(session_id) { - return Err(Error::invalid_params().data("session already has an active prompt")); - } - Ok(()) - } - - fn begin_prompt(&mut self, session_id: &str) -> Result { - self.ensure_attachable(session_id)?; - let session = self - .attachments - .get_mut(session_id) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - let cancellation = CancellationToken::new(); - session.cancellation = Some(cancellation.clone()); - self.active_prompts.insert(session_id.to_string()); - Ok(cancellation) - } - - // Call only after the stream and wrapper have finished every persisted write. - fn finish_prompt(&mut self, session_id: &str) { - if let Some(session) = self.attachments.get_mut(session_id) { - session.cancellation = None; - } - self.active_prompts.remove(session_id); - } -} - -fn permission_options() -> Vec { - vec![ - PermissionOption::new("once", "Allow once", PermissionOptionKind::AllowOnce), - PermissionOption::new( - "always", - "Always allow for this session", - PermissionOptionKind::AllowAlways, - ), - PermissionOption::new("reject", "Reject", PermissionOptionKind::RejectOnce), - ] -} - -fn model_output_limit(config: &LoadedConfig, provider_id: &str, model_id: &str) -> Option { - if let Some(limit) = config - .merged_config - .custom_providers - .get(&provider_id.trim().to_ascii_lowercase()) - .and_then(|provider| provider.models.get(model_id)) - .and_then(|model| model.max_tokens) - { - return Some(limit); - } - crate::model::discovery::Discovery::new_with_config( - Some(config.merged_config.custom_providers.clone()), - config.merged_config.disabled_providers.clone(), - config.merged_config.enabled_providers.clone(), - ) - .ok() - .and_then(|discovery| discovery.get_model_output_limit(provider_id, model_id)) -} - -fn command_text(parts: &[ContentBlock]) -> String { - parts - .iter() - .find_map(|part| match part { - ContentBlock::Text(content) => Some(content.text.clone()), - _ => None, - }) - .unwrap_or_default() -} - -fn acp_session_info( - session: crate::session::manager::SessionInfo, - root_id: Option, -) -> SessionInfo { - let mut meta = serde_json::Map::new(); - meta.insert( - "crabcode".to_string(), - serde_json::json!({ - "parentSessionId": session.parent_id, - "rootSessionId": root_id, - }), - ); - SessionInfo::new(session.id, session.workspace_path) - .title(session.title) - .updated_at(system_time_to_iso8601(session.updated_at)) - .meta(meta) -} - -fn session_page( - cursor: Option<&str>, - total: usize, -) -> Result<(usize, usize, Option), Error> { - let offset = cursor - .unwrap_or("0") - .parse::() - .map_err(|_| Error::invalid_params().data("invalid session list cursor"))?; - let end = offset.saturating_add(100).min(total); - let next_cursor = (end < total).then(|| end.to_string()); - Ok((offset.min(total), end, next_cursor)) -} - -fn mcp_native_tool_content(metadata: &serde_json::Value) -> Vec { - metadata - .get("mcp_result") - .and_then(|result| result.get("content")) - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter(|block| { - !matches!( - block.get("type").and_then(serde_json::Value::as_str), - Some("text") - ) - }) - .map(|block| { - serde_json::from_value::(block.clone()) - .map(ToolCallContent::from) - .unwrap_or_else(|_| { - let kind = block - .get("type") - .and_then(serde_json::Value::as_str) - .unwrap_or("unknown"); - ToolCallContent::from(format!( - "[MCP {kind} content]\n{}", - serde_json::to_string_pretty(block).unwrap_or_else(|_| block.to_string()) - )) - }) - }) - .collect() -} - -fn metadata_has_mcp_images(metadata: Option<&serde_json::Value>) -> bool { - metadata - .and_then(|metadata| metadata.get("mcp_result")) - .and_then(|result| result.get("content")) - .and_then(serde_json::Value::as_array) - .is_some_and(|content| { - content - .iter() - .any(|block| block.get("type").and_then(serde_json::Value::as_str) == Some("image")) - }) -} - -fn permission_tool_content( - prompt: &crate::tools::PermissionPrompt, - cwd: &Path, -) -> Vec { - match prompt.tool_id.as_str() { - "write" => prompt - .raw_input - .get("file_path") - .or_else(|| prompt.raw_input.get("filePath")) - .and_then(serde_json::Value::as_str) - .zip( - prompt - .raw_input - .get("content") - .and_then(serde_json::Value::as_str), - ) - .map(|(path, new_text)| vec![preflight_diff(path, new_text, cwd)]) - .unwrap_or_default(), - "write_files" => prompt - .raw_input - .get("files") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(|file| { - let path = file.get("file_path")?.as_str()?; - let content = file.get("content")?.as_str()?; - Some(preflight_diff(path, content, cwd)) - }) - .collect(), - "edit" => permission_edit_diff(&prompt.raw_input, cwd) - .into_iter() - .collect(), - "apply_patch" - if crate::tools::patch::patch_paths_from_params(&prompt.raw_input) - .iter() - .all(|path| preview_read_allowed(&absolute_tool_path(path, cwd), cwd)) => - { - crate::tools::patch::preview_patch(&prompt.raw_input, cwd) - .unwrap_or_default() - .into_iter() - .map(|change| { - agent_client_protocol::schema::v1::Diff::new(change.path, change.new_text) - .old_text(change.old_text) - .into() - }) - .collect() - } - _ => Vec::new(), - } -} - -// A write/edit approval is not a read approval. Keep ordinary workspace previews, -// but never acquire protected file contents just to display a permission request. -fn preview_read_allowed(path: &Path, cwd: &Path) -> bool { - !crate::tools::permission::is_outside_workdir(path, cwd) - && !crate::tools::permission::is_sensitive_path(path) - && !std::fs::canonicalize(path) - .ok() - .is_some_and(|resolved| crate::tools::permission::is_sensitive_path(&resolved)) -} - -async fn cancellable( - token: &CancellationToken, - future: impl std::future::Future, -) -> Option { - tokio::select! { - biased; - _ = token.cancelled() => None, - result = future => Some(result), - } -} - -fn preflight_diff(path: &str, new_text: &str, cwd: &Path) -> ToolCallContent { - let path = absolute_tool_path(path, cwd); - let old_text = preview_read_allowed(&path, cwd) - .then(|| std::fs::read_to_string(&path).ok()) - .flatten(); - agent_client_protocol::schema::v1::Diff::new(path, new_text.to_string()) - .old_text(old_text) - .into() -} - -fn permission_edit_diff(input: &serde_json::Value, cwd: &Path) -> Option { - let path = input - .get("file_path") - .or_else(|| input.get("filePath"))? - .as_str()?; - let old_string = input.get("old_string")?.as_str()?; - let new_string = input.get("new_string")?.as_str()?; - let absolute = absolute_tool_path(path, cwd); - if !preview_read_allowed(&absolute, cwd) { - return None; - } - let old_text = std::fs::read_to_string(&absolute).ok()?; - let new_text = if input - .get("replace_all") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - { - old_text.replace(old_string, new_string) - } else { - old_text.replacen(old_string, new_string, 1) - }; - Some( - agent_client_protocol::schema::v1::Diff::new(absolute, new_text) - .old_text(old_text) - .into(), - ) -} - -fn write_prompt_audio( - session_id: &str, - audio: &agent_client_protocol::schema::v1::AudioContent, -) -> Result { - const MAX_AUDIO_BYTES: usize = 20 * 1024 * 1024; - let media_type = audio.mime_type.trim().to_ascii_lowercase(); - let extension = match media_type.as_str() { - "audio/wav" | "audio/x-wav" | "audio/wave" => "wav", - "audio/mpeg" | "audio/mp3" => "mp3", - _ => { - return Err(Error::invalid_params() - .data(format!("unsupported audio MIME type: {}", audio.mime_type))); - } - }; - let data = base64::engine::general_purpose::STANDARD - .decode(audio.data.trim()) - .map_err(|error| Error::invalid_params().data(format!("invalid audio data: {error}")))?; - if data.len() > MAX_AUDIO_BYTES { - return Err(Error::invalid_params().data("audio exceeds the 20 MiB size limit")); - } - let path = crate::persistence::attachments::write(session_id, extension, &data) - .map_err(|_| internal_error())?; - Ok(path.to_string_lossy().into_owned()) -} - -fn model_supports_audio(config: &LoadedConfig, provider: &str, model: &str) -> bool { - crate::model::discovery::Discovery::new_with_custom(Some( - config.merged_config.custom_providers.clone(), - )) - .ok() - .is_some_and(|discovery| discovery.model_supports_input_modality(provider, model, "audio")) -} - -struct ManagedAttachmentGuard { - paths: Vec, - committed: bool, -} - -impl ManagedAttachmentGuard { - fn new(paths: Vec) -> Self { - Self { - paths, - committed: false, - } - } - - fn commit(&mut self) { - self.committed = true; - } -} - -impl Drop for ManagedAttachmentGuard { - fn drop(&mut self) { - if !self.committed { - for path in &self.paths { - crate::persistence::attachments::remove_file(Path::new(path)); - } - } - } -} - -struct AcpQuestionField { - selection: String, - custom: String, - labels: HashMap, - multiple: bool, -} - -struct AcpQuestionForm { - request: CreateElicitationRequest, - fields: Vec, -} - -fn skipped_question_answers(questions: &serde_json::Value) -> serde_json::Value { - let count = questions.as_array().map_or(1, Vec::len); - serde_json::Value::Array( - (0..count) - .map(|_| serde_json::Value::Array(Vec::new())) - .collect(), - ) -} - -fn question_text(question: &serde_json::Value, key: &str, fallback: &str) -> String { - question - .get(key) - .and_then(serde_json::Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or(fallback) - .to_string() -} - -fn acp_question_form( - session_id: &str, - tool_call_id: Option<&str>, - questions: &serde_json::Value, -) -> AcpQuestionForm { - let question_items = questions - .as_array() - .cloned() - .unwrap_or_else(|| vec![questions.clone()]); - let mut schema = ElicitationSchema::new() - .title("Agent questions") - .description("Answer any fields you want; blank fields are treated as skipped."); - let mut fields = Vec::with_capacity(question_items.len()); - - for (question_index, question) in question_items.iter().enumerate() { - let selection = format!("question_{question_index}"); - let custom = format!("question_{question_index}_custom"); - let prompt = question_text(question, "question", "Question"); - let header = question_text( - question, - "header", - &format!("Question {}", question_index + 1), - ); - let mut labels = HashMap::new(); - let options = question - .get("options") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .enumerate() - .filter_map(|(option_index, option)| { - let label = option - .get("label") - .and_then(serde_json::Value::as_str) - .or_else(|| option.as_str())? - .trim(); - if label.is_empty() { - return None; - } - - let value = format!("q{question_index}_option_{option_index}"); - labels.insert(value.clone(), label.to_string()); - let description = option - .get("description") - .and_then(serde_json::Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()); - let mut option = EnumOption::new(value, label); - if let Some(description) = description { - option = option.description(description); - } - Some(option) - }) - .collect::>(); - let multiple = question - .get("multiple") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false); - - if multiple { - schema = schema.property( - selection.clone(), - MultiSelectPropertySchema::titled(options) - .title(header.clone()) - .description(prompt.clone()), - false, - ); - } else { - schema = schema.property( - selection.clone(), - StringPropertySchema::new() - .title(header.clone()) - .description(prompt.clone()) - .one_of(options), - false, - ); - } - schema = schema.property( - custom.clone(), - StringPropertySchema::new() - .title(format!("{header}: custom answer")) - .description("Optional free-form answer."), - false, - ); - fields.push(AcpQuestionField { - selection, - custom, - labels, - multiple, - }); - } - - let scope = ElicitationSessionScope::new(session_id.to_string()) - .tool_call_id(tool_call_id.map(agent_client_protocol::schema::v1::ToolCallId::new)); - let request = CreateElicitationRequest::new( - ElicitationFormMode::new(scope, schema), - "The agent needs additional input to continue.", - ); - AcpQuestionForm { request, fields } -} - -fn acp_question_answers( - fields: &[AcpQuestionField], - action: ElicitationAction, -) -> serde_json::Value { - let ElicitationAction::Accept(accepted) = action else { - return serde_json::Value::Array( - fields - .iter() - .map(|_| serde_json::Value::Array(Vec::new())) - .collect(), - ); - }; - let content = accepted.content.unwrap_or_default(); - serde_json::Value::Array( - fields - .iter() - .map(|field| { - let mut answers = Vec::new(); - match content.get(&field.selection) { - Some(ElicitationContentValue::String(value)) if !field.multiple => { - if let Some(label) = field.labels.get(value) { - answers.push(serde_json::Value::String(label.clone())); - } - } - Some(ElicitationContentValue::StringArray(values)) if field.multiple => { - for value in values { - if let Some(label) = field.labels.get(value) { - if !answers.iter().any(|answer| answer.as_str() == Some(label)) { - answers.push(serde_json::Value::String(label.clone())); - } - } - } - } - _ => {} - } - if let Some(ElicitationContentValue::String(custom)) = content.get(&field.custom) { - let custom = custom.trim(); - if !custom.is_empty() { - if !field.multiple { - answers.clear(); - } - answers.push(serde_json::Value::String( - custom.chars().take(8_192).collect(), - )); - } - } - serde_json::Value::Array(answers) - }) - .collect(), - ) -} - -fn compact_command(prompt: &str) -> Result { - let trimmed = prompt.trim(); - let Some(command_line) = trimmed.strip_prefix('/') else { - return Ok(false); - }; - let (name, args) = command_line - .split_once(char::is_whitespace) - .map(|(name, args)| (name, args.trim())) - .unwrap_or((command_line, "")); - if name != "compact" { - return Ok(false); - } - if !args.is_empty() { - return Err(Error::invalid_params().data("Usage: /compact")); - } - Ok(true) -} - -fn permission_tool_call_id(tool_call_id: Option<&str>) -> String { - tool_call_id - .map(str::to_string) - .unwrap_or_else(|| format!("permission:{}", cuid2::create_id())) -} - -fn resolved_reasoning( - session: &AcpSession, - requested: crate::model::reasoning::ReasoningEffort, -) -> Option { - model_reasoning_capability( - &session.config, - &session.models, - &session.provider, - &session.model, - ) - .and_then(|capability| capability.resolve(Some(requested))) -} - -fn reasoning_effort_label(effort: crate::model::reasoning::ReasoningEffort) -> &'static str { - match effort { - crate::model::reasoning::ReasoningEffort::None => "None", - crate::model::reasoning::ReasoningEffort::Minimal => "Minimal", - crate::model::reasoning::ReasoningEffort::Low => "Low", - crate::model::reasoning::ReasoningEffort::Medium => "Medium", - crate::model::reasoning::ReasoningEffort::High => "High", - crate::model::reasoning::ReasoningEffort::XHigh => "Xhigh", - crate::model::reasoning::ReasoningEffort::Max => "Max", - } -} - -fn available_commands(session: &AcpSession) -> Vec { - let mut commands: Vec<_> = session - .config - .merged_config - .commands - .iter() - .filter(|command| !matches!(command.name.as_str(), "btw" | "compact" | "skills" | "mcp")) - .map(|command| { - let description = command - .description - .clone() - .unwrap_or_else(|| format!("Run /{}", command.name)); - AvailableCommand::new(command.name.clone(), description).input( - AvailableCommandInput::Unstructured(UnstructuredCommandInput::new("Arguments")), - ) - }) - .collect(); - commands.extend( - session - .skills - .all() - .into_iter() - .filter(|skill| !matches!(skill.name.as_str(), "btw" | "compact" | "skills" | "mcp")) - .map(|skill| { - AvailableCommand::new( - skill.name.clone(), - skill - .description - .clone() - .unwrap_or_else(|| format!("Use the {} skill", skill.name)), - ) - .input(AvailableCommandInput::Unstructured( - UnstructuredCommandInput::new("Task or context for this skill"), - )) - }), - ); - commands.push(AvailableCommand::new( - "skills", - "List skills available in this workspace", - )); - commands.push(AvailableCommand::new( - "mcp", - "List configured MCP servers and their status", - )); - commands.push(AvailableCommand::new( - "compact", - "Summarize this session to reduce context", - )); - commands.push( - AvailableCommand::new( - "btw", - "Ask a quick side question without changing session history", - ) - .input(AvailableCommandInput::Unstructured( - UnstructuredCommandInput::new("Question"), - )), - ); - commands.sort_by(|left, right| left.name.cmp(&right.name)); - commands.dedup_by(|left, right| left.name == right.name); - commands -} - -#[derive(Debug, PartialEq, Eq)] -enum SlashExpansion { - Prompt { - prompt: String, - agent: Option, - model: Option, - }, - LocalResult(String), - Btw(String), -} - -async fn expand_slash_command(session: &AcpSession, prompt: &str) -> Result { - let Some(command_line) = prompt.strip_prefix('/') else { - return Ok(SlashExpansion::Prompt { - prompt: prompt.to_string(), - agent: None, - model: None, - }); - }; - let (name, args) = command_line - .split_once(char::is_whitespace) - .map(|(name, args)| (name, args.trim_start())) - .unwrap_or((command_line, "")); - if name == "btw" { - let question = args.trim(); - if question.is_empty() { - return Err(Error::invalid_params().data("Usage: /btw ")); - } - return Ok(SlashExpansion::Btw(question.to_string())); - } - if name == "skills" { - if !args.is_empty() { - return Err(Error::invalid_params().data("Usage: /skills")); - } - let skills = session - .skills - .all() - .into_iter() - .map(|skill| { - format!( - "- /{} — {}", - skill.name, - skill.description.as_deref().unwrap_or("No description") - ) - }) - .collect::>(); - return Ok(SlashExpansion::LocalResult(if skills.is_empty() { - "No skills are available in this workspace.".to_string() - } else { - format!("Available skills:\n{}", skills.join("\n")) - })); - } - if name == "mcp" { - if !args.is_empty() { - return Err(Error::invalid_params().data("Usage: /mcp")); - } - let manager = crate::mcp::McpManager::ensure( - session.config.merged_config.mcp.clone(), - session.cwd.clone(), - ); - let servers = manager - .lock() - .await - .views() - .into_iter() - .map(|server| { - let detail = server - .detail - .map(|detail| format!(": {detail}")) - .unwrap_or_default(); - format!( - "- {} ({}, {}){}", - server.name, server.kind, server.status, detail - ) - }) - .collect::>(); - return Ok(SlashExpansion::LocalResult(if servers.is_empty() { - "No MCP servers are configured for this workspace.".to_string() - } else { - format!("MCP servers:\n{}", servers.join("\n")) - })); - } - if let Some(command) = session - .config - .merged_config - .commands - .iter() - .find(|command| command.name == name) - { - return command - .render(args) - .await - .map_err(|_| internal_error()) - .map(|rendered| SlashExpansion::Prompt { - prompt: rendered.prompt, - agent: rendered.agent, - model: rendered.model, - }); - } - if let Some(skill) = session.skills.get(name) { - let mut expanded = skill.content.clone(); - if !args.is_empty() { - expanded.push_str("\n\nUser task/context:\n"); - expanded.push_str(args); - } - return Ok(SlashExpansion::Prompt { - prompt: expanded, - agent: None, - model: None, - }); - } - Err(Error::invalid_params().data(format!("Unknown ACP command: /{name}"))) -} - -fn merge_acp_mcp_servers(config: &mut LoadedConfig, servers: Vec) { - for server in servers { - let (name, server) = match server { - McpServer::Stdio(server) => { - let mut command = vec![server.command.to_string_lossy().into_owned()]; - command.extend(server.args); - ( - server.name, - crate::config::configuration::McpServerConfig::Local( - crate::config::configuration::McpLocalConfig { - command, - cwd: None, - environment: server - .env - .into_iter() - .map(|variable| (variable.name, variable.value)) - .collect(), - enabled: true, - timeout_ms: None, - }, - ), - ) - } - McpServer::Http(server) => ( - server.name, - crate::config::configuration::McpServerConfig::Remote( - crate::config::configuration::McpRemoteConfig { - url: server.url, - headers: server - .headers - .into_iter() - .map(|header| (header.name, header.value)) - .collect(), - enabled: true, - timeout_ms: None, - oauth_enabled: false, - oauth_client_id: None, - oauth_client_secret: None, - oauth_scope: None, - }, - ), - ), - McpServer::Sse(server) => ( - server.name, - crate::config::configuration::McpServerConfig::Remote( - crate::config::configuration::McpRemoteConfig { - url: server.url, - headers: server - .headers - .into_iter() - .map(|header| (header.name, header.value)) - .collect(), - enabled: true, - timeout_ms: None, - oauth_enabled: false, - oauth_client_id: None, - oauth_client_secret: None, - oauth_scope: None, - }, - ), - ), - #[allow(unreachable_patterns)] - _ => continue, - }; - config.merged_config.mcp.insert(name, server); - } -} - -#[derive(Clone)] -struct AcpSession { - cwd: PathBuf, - config: LoadedConfig, - tool_permissions: crate::tools::ToolPermissions, - skills: crate::skill::SkillStore, - models: Vec, - provider: String, - model: String, - agent: String, - reasoning_selection: crate::model::reasoning::ReasoningEffort, - reasoning: Option, - context_window: Option, - cancellation: Option, -} - -impl AcpService { - pub fn new(initial_workspace: &Path) -> Result { - let session_manager = SessionManager::new() - .with_history_for_workspace(initial_workspace) - .map_err(|_| internal_error())?; - Ok(Self { - sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), - session_manager: Arc::new(Mutex::new(session_manager)), - client_capabilities: Arc::new(Mutex::new(Default::default())), - }) - } - - pub fn set_client_capabilities( - &self, - capabilities: agent_client_protocol::schema::v1::ClientCapabilities, - ) { - if let Ok(mut current) = self.client_capabilities.lock() { - *current = capabilities; - } - } - - fn supports_form_elicitation(&self) -> bool { - self.client_capabilities - .lock() - .ok() - .and_then(|capabilities| capabilities.elicitation.clone()) - .and_then(|elicitation| elicitation.form) - .is_some() - } - - fn supports_terminals(&self) -> bool { - self.client_capabilities - .lock() - .ok() - .is_some_and(|capabilities| capabilities.terminal) - } - - pub async fn available_commands( - &self, - session_id: &str, - ) -> Result { - let sessions = self.sessions.lock().await; - let session = sessions - .attachments - .get(session_id) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - Ok(AvailableCommandsUpdate::new(available_commands(session))) - } - - pub async fn new_session( - &self, - cwd: PathBuf, - mcp_servers: Vec, - ) -> Result { - let cwd = workspace_path(&cwd)?; - let mut config = - crate::config::ConfigLoader::load_for(&cwd).map_err(|_| internal_error())?; - merge_acp_mcp_servers(&mut config, mcp_servers); - crate::skill::init_skill_store(&config.xdg_config_home, &config.project_root); - let (provider, model) = resolve_model(&config); - let models = crate::model::catalog::selectable_models(&config, None) - .await - .map_err(|_| internal_error())?; - let reasoning = model_reasoning(&config, &models, &provider, &model); - let reasoning_selection = - reasoning.unwrap_or(crate::model::reasoning::ReasoningEffort::None); - let context_window = model_context_window(&config, &models, &provider, &model); - let agent = config - .merged_config - .default_agent - .clone() - .unwrap_or_else(|| { - config - .merged_config - .agent_registry - .default_agent() - .to_string() - }); - - let session_id = { - let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; - manager - .switch_current_workspace_path(&cwd.to_string_lossy()) - .map_err(|_| internal_error())?; - manager.create_session(Some("New session".to_string())) - }; - - let skills = crate::skill::SkillStore::load(&config.xdg_config_home, &config.project_root); - let tool_permissions = configured_tool_permissions(&cwd, &config); - self.sessions.lock().await.attachments.insert( - session_id.clone(), - AcpSession { - cwd, - config, - tool_permissions, - skills, - models, - provider, - model, - agent, - reasoning_selection, - reasoning, - context_window, - cancellation: None, - }, - ); - - let session = self - .sessions - .lock() - .await - .attachments - .get(&session_id) - .cloned() - .ok_or_else(internal_error)?; - Ok(NewSessionResponse::new(session_id) - .modes(session_modes(&session)) - .config_options(session_config_options(&session))) - } - - pub async fn list_sessions( - &self, - cwd: Option, - cursor: Option, - ) -> Result { - let cwd = cwd.as_deref().map(workspace_path).transpose()?; - let manager = self.session_manager.lock().map_err(|_| internal_error())?; - let mut sessions = manager - .list_sessions() - .into_iter() - .filter(|session| session.archived_at.is_none()) - .filter(|session| { - cwd.as_ref() - .is_none_or(|cwd| session.workspace_path == cwd.to_string_lossy()) - }) - .collect::>(); - sessions.sort_by(|left, right| right.updated_at.cmp(&left.updated_at)); - - let (offset, end, next_cursor) = session_page(cursor.as_deref(), sessions.len())?; - Ok(ListSessionsResponse::new( - sessions - .into_iter() - .skip(offset) - .take(end.saturating_sub(offset)) - .map(|session| { - let root_id = manager.root_session_id_for(&session.id); - acp_session_info(session, root_id) - }) - .collect(), - ) - .next_cursor(next_cursor)) - } - - pub async fn load_session( - &self, - session_id: String, - cwd: PathBuf, - connection: ConnectionTo, - ) -> Result { - let (session, messages) = self.attach_persisted_session(&session_id, cwd).await?; - replay_messages(&connection, &session_id, &messages, &session.cwd)?; - Ok(LoadSessionResponse::new() - .modes(session_modes(&session)) - .config_options(session_config_options(&session))) - } - - pub async fn resume_session( - &self, - session_id: String, - cwd: PathBuf, - ) -> Result { - let (session, _) = self.attach_persisted_session(&session_id, cwd).await?; - Ok(ResumeSessionResponse::new() - .modes(session_modes(&session)) - .config_options(session_config_options(&session))) - } - - pub async fn fork_session( - &self, - session_id: String, - cwd: PathBuf, - ) -> Result { - let (source, messages) = self.attach_persisted_session(&session_id, cwd).await?; - let fork_id = { - let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; - let source_title = manager - .get_session(&session_id) - .map(|session| session.title.clone()) - .unwrap_or_else(|| session_id.clone()); - manager - .switch_current_workspace_path(&source.cwd.to_string_lossy()) - .map_err(|_| internal_error())?; - let fork_id = manager.create_session(Some(format!("{source_title} (fork)"))); - let messages = - match crate::persistence::attachments::clone_messages(&messages, &fork_id) { - Ok(messages) => messages, - Err(_) => { - manager.delete_session(&fork_id); - return Err(internal_error()); - } - }; - if manager - .replace_session_messages(&fork_id, messages) - .is_err() - { - manager.delete_session(&fork_id); - return Err(internal_error()); - } - fork_id - }; - self.sessions.lock().await.attachments.insert( - fork_id.clone(), - AcpSession { - tool_permissions: configured_tool_permissions(&source.cwd, &source.config), - ..source.clone() - }, - ); - Ok(NewSessionResponse::new(fork_id) - .modes(session_modes(&source)) - .config_options(session_config_options(&source))) - } - - pub async fn close_session(&self, session_id: &str) { - if let Some(session) = self.sessions.lock().await.attachments.remove(session_id) { - if let Some(cancellation) = session.cancellation { - cancellation.cancel(); - } - } - } - - pub async fn delete_session(&self, session_id: &str) -> Result<(), Error> { - self.close_session(session_id).await; - let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; - match manager.try_delete_session(session_id) { - Ok(true) => Ok(()), - Ok(false) => Err(Error::invalid_params().data("unknown session")), - Err(_) => Err(internal_error()), - } - } - - pub async fn cancel_session(&self, session_id: &str) { - if let Some(session) = self.sessions.lock().await.attachments.get(session_id) { - if let Some(cancellation) = &session.cancellation { - cancellation.cancel(); - } - } - } - - pub async fn set_mode( - &self, - session_id: &str, - mode_id: &str, - ) -> Result { - let mut sessions = self.sessions.lock().await; - let session = sessions - .attachments - .get_mut(session_id) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - if session - .config - .merged_config - .agent_registry - .primary_agent(mode_id) - .is_none_or(|agent| agent.hidden) - { - return Err(Error::invalid_params().data("unknown ACP mode")); - } - session.agent = mode_id.to_string(); - Ok(SetSessionConfigOptionResponse::new(session_config_options( - session, - ))) - } - - pub async fn set_model( - &self, - session_id: &str, - model_ref: &str, - ) -> Result { - let mut sessions = self.sessions.lock().await; - let session = sessions - .attachments - .get_mut(session_id) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - let model = find_selectable_model(&session.models, model_ref)?; - session.provider.clone_from(&model.provider_id); - session.model.clone_from(&model.id); - session.reasoning = resolved_reasoning(session, session.reasoning_selection); - session.context_window = model_context_window( - &session.config, - &session.models, - &session.provider, - &session.model, - ); - Ok(SetSessionConfigOptionResponse::new(session_config_options( - session, - ))) - } - - pub async fn set_reasoning_effort( - &self, - session_id: &str, - value: &str, - ) -> Result { - let mut sessions = self.sessions.lock().await; - let session = sessions - .attachments - .get_mut(session_id) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - let requested = value - .parse::() - .map_err(|_| Error::invalid_params().data("invalid reasoning effort"))?; - session.reasoning_selection = requested; - session.reasoning = resolved_reasoning(session, requested); - Ok(SetSessionConfigOptionResponse::new(session_config_options( - session, - ))) - } - - async fn attach_persisted_session( - &self, - session_id: &str, - cwd: PathBuf, - ) -> Result<(AcpSession, Vec), Error> { - // Fail promptly during close cleanup, before model/config discovery. Recheck - // below after awaits and hold the lock through the actual attachment. - self.sessions.lock().await.ensure_attachable(session_id)?; - let cwd = workspace_path(&cwd)?; - // Invalid load/resume requests must not wait on network model discovery. - { - let manager = self.session_manager.lock().map_err(|_| internal_error())?; - let stored = manager - .get_session_ref(session_id) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - if stored.workspace_path != cwd.to_string_lossy() { - return Err(Error::invalid_params().data("session does not belong to cwd")); - } - } - let config = crate::config::ConfigLoader::load_for(&cwd).map_err(|_| internal_error())?; - crate::skill::init_skill_store(&config.xdg_config_home, &config.project_root); - let models = crate::model::catalog::selectable_models(&config, None) - .await - .map_err(|_| internal_error())?; - - let mut sessions = self.sessions.lock().await; - sessions.ensure_attachable(session_id)?; - let messages = { - let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; - let stored = manager - .get_session(session_id) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - if stored.workspace_path != cwd.to_string_lossy() { - return Err(Error::invalid_params().data("session does not belong to cwd")); - } - stored.messages.clone() - }; - - let (provider, model) = messages - .iter() - .rev() - .find_map(|message| match (&message.provider, &message.model) { - (Some(provider), Some(model)) => Some((provider.clone(), model.clone())), - _ => None, - }) - .unwrap_or_else(|| resolve_model(&config)); - let agent = messages - .iter() - .rev() - .find_map(|message| message.agent_mode.clone()) - .unwrap_or_else(|| { - config - .merged_config - .default_agent - .clone() - .unwrap_or_else(|| { - config - .merged_config - .agent_registry - .default_agent() - .to_string() - }) - }); - let reasoning = model_reasoning(&config, &models, &provider, &model); - let reasoning_selection = - reasoning.unwrap_or(crate::model::reasoning::ReasoningEffort::None); - let context_window = model_context_window(&config, &models, &provider, &model); - let skills = crate::skill::SkillStore::load(&config.xdg_config_home, &config.project_root); - let tool_permissions = configured_tool_permissions(&cwd, &config); - let session = AcpSession { - cwd, - config, - tool_permissions, - skills, - models, - provider, - model, - agent, - reasoning_selection, - reasoning, - context_window, - cancellation: None, - }; - sessions - .attachments - .insert(session_id.to_string(), session.clone()); - Ok((session, messages)) - } - - pub async fn prompt( - &self, - session_id: String, - prompt: Vec, - connection: ConnectionTo, - ) -> Result { - let cancellation = self.sessions.lock().await.begin_prompt(&session_id)?; - // Only individual blocking operations are cancellable. Dropping prompt_inner - // would skip transcript persistence and leave the session marked active. - let mut result = self - .prompt_inner(session_id.clone(), prompt, connection, cancellation.clone()) - .await; - if cancellation.is_cancelled() { - result = Ok(PromptResponse::new(StopReason::Cancelled)); - } - // Cover setup errors and local commands as well as stream/compaction exits. - if let Ok(mut manager) = self.session_manager.lock() { - let status = if cancellation.is_cancelled() - || result - .as_ref() - .is_ok_and(|response| response.stop_reason == StopReason::Cancelled) - { - crate::session::types::SessionStatus::Interrupted - } else if result.is_err() { - crate::session::types::SessionStatus::Failed - } else { - crate::session::types::SessionStatus::Idle - }; - if manager.get_session_ref(&session_id).is_some() { - if manager - .set_session_status(&session_id, status, None) - .is_err() - { - result = Err(internal_error()); - } - } - } else { - result = Err(internal_error()); - } - // Also stop any detached producer after an error or early return. - cancellation.cancel(); - self.sessions.lock().await.finish_prompt(&session_id); - result - } - - async fn prompt_inner( - &self, - session_id: String, - prompt: Vec, - connection: ConnectionTo, - cancellation: CancellationToken, - ) -> Result { - let mut session = self - .sessions - .lock() - .await - .attachments - .get(&session_id) - .cloned() - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - let compact_text = command_text(&prompt); - if compact_command(&compact_text)? { - if prompt - .iter() - .any(|part| !matches!(part, ContentBlock::Text(_))) - { - return Err(Error::invalid_params().data("/compact does not accept attachments")); - } - return self - .compact_session(&session_id, session, connection, cancellation) - .await; - } - let expansion = if compact_text.trim_start().starts_with('/') { - let Some(expansion) = - cancellable(&cancellation, expand_slash_command(&session, &compact_text)).await - else { - return Ok(PromptResponse::new(StopReason::Cancelled)); - }; - Some(expansion?) - } else { - None - }; - let expanded_prompt = match expansion { - Some(SlashExpansion::LocalResult(text)) => { - if prompt - .iter() - .any(|part| !matches!(part, ContentBlock::Text(_))) - { - return Err(Error::invalid_params() - .data("local ACP commands do not accept attachments")); - } - let message_id = cuid2::create_id(); - send_replay_text(&connection, &session_id, &message_id, &text, false, false)?; - return Ok(PromptResponse::new(StopReason::EndTurn)); - } - Some(SlashExpansion::Prompt { - prompt, - agent, - model, - }) => { - if let Some(agent) = agent { - if session - .config - .merged_config - .agent_registry - .get(&agent) - .is_none() - { - return Err(Error::invalid_params() - .data(format!("custom command references unknown agent: {agent}"))); - } - session.agent = agent; - } - if let Some(model_ref) = model { - let (provider, model) = crate::app::parse_model_ref(&model_ref); - let canonical = format!("{provider}/{model}"); - let model = find_selectable_model(&session.models, &canonical)?; - session.provider.clone_from(&model.provider_id); - session.model.clone_from(&model.id); - session.reasoning = resolved_reasoning(&session, session.reasoning_selection); - session.context_window = model_context_window( - &session.config, - &session.models, - &session.provider, - &session.model, - ); - } - Some(prompt) - } - Some(SlashExpansion::Btw(question)) => { - if prompt - .iter() - .any(|part| !matches!(part, ContentBlock::Text(_))) - { - return Err(Error::invalid_params().data("/btw does not accept attachments")); - } - return self - .btw_session(&session_id, session, question, connection, cancellation) - .await; - } - None => None, - }; - let supports_images = session - .models - .iter() - .find(|model| model.provider_id == session.provider && model.id == session.model) - .is_some_and(|model| model.attachment); - let supports_audio = - model_supports_audio(&session.config, &session.provider, &session.model); - let (prompt, local_image_paths, local_audio_paths) = prompt_content( - prompt, - expanded_prompt.as_deref(), - supports_images, - supports_audio, - &session_id, - &session, - )?; - let mut managed_paths = local_image_paths.clone(); - managed_paths.extend(local_audio_paths.clone()); - let mut attachment_guard = ManagedAttachmentGuard::new(managed_paths); - if prompt.trim().is_empty() { - return Err(Error::invalid_params().data("prompt must include text content")); - } - let mut messages = { - let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; - let stored = manager - .get_session(&session_id) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - stored.messages.clone() - }; - let mut user_message = crate::session::types::Message::user(&prompt); - user_message.local_image_paths = local_image_paths; - user_message.local_audio_paths = local_audio_paths; - user_message.provider = Some(session.provider.clone()); - user_message.model = Some(session.model.clone()); - user_message.agent_mode = Some(session.agent.clone()); - let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); - let process_registry = std::sync::Arc::new(crate::tools::ProcessRegistry::new()); - let prompt_registry = cancellable( - &cancellation, - crate::tools::initialize_tool_registry_with_dynamic_config( - Some(sender.clone()), - tool_permissions(&session), - session.config.merged_config.agent_registry.clone(), - cancellation.clone(), - Some(&session.provider), - &session.config.merged_config.websearch, - &session.config.merged_config.mcp, - &session.cwd, - process_registry.clone(), - ), - ) - .await; - let Some(prompt_registry) = prompt_registry else { - return Ok(PromptResponse::new(StopReason::Cancelled)); - }; - let is_git_repo = - crate::utils::git::is_git_repo(&session.cwd.to_string_lossy()).unwrap_or(false); - let system_prompt = cancellable( - &cancellation, - crate::prompt::SystemPromptComposer::new( - &session.model, - session.cwd.to_string_lossy(), - is_git_repo, - std::env::consts::OS, - ) - .with_tool_registry(prompt_registry.clone()) - .with_agent_registry(session.config.merged_config.agent_registry.clone()) - .with_active_agent(session.agent.clone()) - .compose(), - ) - .await; - let Some(system_prompt) = system_prompt else { - return Ok(PromptResponse::new(StopReason::Cancelled)); - }; - let system_message = crate::session::types::Message::system(system_prompt); - let prompt_tokens = crate::session::compaction::message_context_tokens(&system_message); - let auto_compacted = self - .maybe_auto_compact_session( - &session_id, - &session, - &connection, - cancellation.clone(), - prompt_tokens, - crate::session::compaction::message_context_tokens(&user_message), - ) - .await?; - if cancellation.is_cancelled() { - return Ok(PromptResponse::new(StopReason::Cancelled)); - } - if auto_compacted { - messages = self - .session_manager - .lock() - .map_err(|_| internal_error())? - .get_session_ref(&session_id) - .map(|stored| stored.messages.clone()) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - } - { - let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; - manager - .add_message_to_session(&session_id, &user_message) - .map_err(|_| internal_error())?; - attachment_guard.commit(); - manager - .set_session_status( - &session_id, - crate::session::types::SessionStatus::Streaming, - None, - ) - .map_err(|_| internal_error())?; - } - messages.push(user_message); - let base_context_tokens = acp_context_tokens(&messages, prompt_tokens); - let base_cost = acp_session_cost(&messages); - messages = acp_request_messages(&messages, system_message); - send_usage( - &connection, - &session_id, - &session, - base_context_tokens, - (base_cost > 0.0).then_some(base_cost), - None, - )?; - - let stream_session_id = session_id.clone(); - let stream_cancellation = cancellation.clone(); - let stream_sender = sender.clone(); - let stream_session = session.clone(); - let stream_tool_registry = prompt_registry; - let stream_process_registry = process_registry; - tokio::spawn(async move { - let result = crate::llm::client::stream_llm_with_cancellation( - stream_cancellation, - stream_session_id, - stream_session.provider.clone(), - stream_session.model.clone(), - stream_session.reasoning, - stream_session.agent.clone(), - stream_session - .config - .merged_config - .agent_registry - .get(&stream_session.agent) - .and_then(|agent| agent.max_steps), - stream_session.config.merged_config.agent_registry.clone(), - tool_permissions(&stream_session), - stream_session.config.merged_config.websearch.clone(), - stream_session.config.merged_config.mcp.clone(), - stream_session.config.merged_config.compaction.clone(), - stream_session.cwd.to_string_lossy().to_string(), - Some(stream_tool_registry), - messages, - sender, - stream_process_registry, - ) - .await; - if let Err(error) = result { - let _ = stream_sender.send(crate::llm::ChunkMessage::Failed(error.to_string())); - } - let _ = stream_sender.send(crate::llm::ChunkMessage::End); - }); - - let response = self - .receive_prompt_stream( - &session_id, - &session, - &connection, - &cancellation, - receiver, - base_context_tokens, - prompt_tokens, - base_cost, - ) - .await?; - if !cancellation.is_cancelled() && response.stop_reason != StopReason::Cancelled { - let _ = self - .maybe_auto_compact_session( - &session_id, - &session, - &connection, - cancellation.clone(), - prompt_tokens, - 0, - ) - .await?; - } - Ok(response) - } - - // Keep finalization outside every cancellable await, including client requests. - async fn receive_prompt_stream( - &self, - session_id: &str, - session: &AcpSession, - connection: &ConnectionTo, - cancellation: &CancellationToken, - mut receiver: tokio::sync::mpsc::UnboundedReceiver, - base_context_tokens: usize, - base_prompt_tokens: usize, - base_cost: f64, - ) -> Result { - let mut assistant = crate::session::types::Message::incomplete(""); - let message_id = assistant.id.clone(); - assistant.provider = Some(session.provider.clone()); - assistant.model = Some(session.model.clone()); - assistant.agent_mode = Some(session.agent.clone()); - let mut failed = None; - let mut cancelled = false; - let mut turn_stop_reason = None; - let mut live_usage = AcpLiveUsage::default(); - - // Notification failures must take the same persistence path as cancellation. - let stream_result: Result<(), Error> = async { - while let Some(chunk) = cancellable(cancellation, receiver.recv()).await.flatten() { - match chunk { - crate::llm::ChunkMessage::Text(text) => { - assistant.append(&text); - send_text(connection, session_id, &message_id, text, false)?; - send_live_usage( - connection, - session_id, - session, - base_context_tokens, - base_cost, - &assistant, - &live_usage, - )?; - } - crate::llm::ChunkMessage::Reasoning(text) => { - assistant.append_reasoning(&text); - send_text(connection, session_id, &message_id, text, true)?; - send_live_usage( - connection, - session_id, - session, - base_context_tokens, - base_cost, - &assistant, - &live_usage, - )?; - } - crate::llm::ChunkMessage::ToolCalls(tool_calls) => { - for tool_call in &tool_calls { - record_tool_call(&mut assistant, tool_call); - } - for tool_call in tool_calls { - send_tool_call(connection, session_id, tool_call, &session.cwd)?; - } - send_live_usage( - connection, - session_id, - session, - base_context_tokens, - base_cost, - &assistant, - &live_usage, - )?; - } - crate::llm::ChunkMessage::ToolResult(result) => { - record_tool_result(&mut assistant, &result); - send_tool_result(connection, session_id, result, &session.cwd)?; - send_live_usage( - connection, - session_id, - session, - base_context_tokens, - base_cost, - &assistant, - &live_usage, - )?; - } - crate::llm::ChunkMessage::Metrics { - token_count, - duration_ms, - usage, - cost, - } => { - assistant.token_count = Some(token_count); - assistant.duration_ms = Some(duration_ms); - if let Some(usage) = usage { - assistant.apply_usage(usage, cost); - live_usage.cumulative = usage; - } - live_usage.cost = cost.unwrap_or(live_usage.cost); - send_usage( - connection, - session_id, - session, - live_usage - .context_tokens(base_context_tokens, &assistant) - .max(base_context_tokens.saturating_add(token_count)), - cumulative_cost(base_cost, live_usage.cost), - usage, - )?; - } - crate::llm::ChunkMessage::Usage(usage) => { - live_usage.observe_provider_usage(&assistant, usage); - live_usage.cost += estimate_session_usage_cost(session, &usage); - assistant - .parts - .push(crate::session::types::MessagePart::usage( - usage.input, - usage.output, - usage.cache_read, - usage.cache_write, - estimate_session_usage_cost(session, &usage), - )); - if let Some(part) = assistant.parts.last_mut() { - part.data["acp_prompt_tokens"] = serde_json::json!(base_prompt_tokens); - } - if usage.output > 0 { - assistant.output_tokens = Some( - assistant - .output_tokens - .unwrap_or(0) - .saturating_add(usage.output as usize), - ); - } - send_live_usage( - connection, - session_id, - session, - base_context_tokens, - base_cost, - &assistant, - &live_usage, - )?; - } - crate::llm::ChunkMessage::Cancelled => { - cancelled = true; - break; - } - crate::llm::ChunkMessage::Failed(error) => failed = Some(error), - crate::llm::ChunkMessage::TurnStopReason(reason) => { - turn_stop_reason = Some(reason) - } - crate::llm::ChunkMessage::PermissionRequest(prompt) => { - let response = cancellable( - cancellation, - request_permission(connection, session_id, &prompt), - ) - .await - .unwrap_or(crate::tools::PermissionResponse::Deny); - let _ = prompt.response_tx.send(response); - } - crate::llm::ChunkMessage::QuestionRequest { - tool_call_id, - questions, - response_tx, - } => { - let response = if self.supports_form_elicitation() { - request_questions( - connection, - session_id, - tool_call_id.as_deref(), - &questions, - cancellation, - ) - .await - } else { - skipped_question_answers(&questions) - }; - let _ = response_tx.send(response); - } - crate::llm::ChunkMessage::TerminalSessionRequest(request) => { - if self.supports_terminals() { - bridge_terminal_session( - connection, - session_id, - &session.cwd, - request, - cancellation, - ) - .await; - } else { - let _ = request - .control_tx - .send(crate::tools::TerminalSessionControl::Stop); - } - } - crate::llm::ChunkMessage::End => break, - _ => {} - } - } - - Ok(()) - } - .await; - if stream_result.is_err() && failed.is_none() { - failed = Some("ACP stream notification failed".to_string()); - } - assistant.is_complete = true; - assistant.was_interrupted = cancelled || cancellation.is_cancelled(); - { - let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; - let persist_result = manager.add_message_to_session(session_id, &assistant); - let status = if assistant.was_interrupted { - crate::session::types::SessionStatus::Interrupted - } else if failed.is_some() { - crate::session::types::SessionStatus::Failed - } else { - crate::session::types::SessionStatus::Idle - }; - manager - .set_session_status(session_id, status, failed.as_deref()) - .map_err(|_| internal_error())?; - persist_result.map_err(|_| internal_error())?; - } - if assistant.was_interrupted { - return Ok(PromptResponse::new(StopReason::Cancelled)); - } - stream_result?; - if let Some(error) = failed { - return Err(internal_error_with(&error)); - } - Ok(PromptResponse::new(acp_stop_reason(turn_stop_reason))) - } - - async fn compact_session( - &self, - session_id: &str, - session: AcpSession, - connection: ConnectionTo, - cancellation: CancellationToken, - ) -> Result { - let status_result = self - .session_manager - .lock() - .map_err(|_| internal_error())? - .set_session_status( - session_id, - crate::session::types::SessionStatus::Streaming, - None, - ); - if status_result.is_err() { - return Err(internal_error()); - } - - let result = self - .run_compaction(session_id, &session, cancellation.clone(), 0) - .await; - self.session_manager - .lock() - .map_err(|_| internal_error())? - .set_session_status( - session_id, - if cancellation.is_cancelled() { - crate::session::types::SessionStatus::Interrupted - } else { - crate::session::types::SessionStatus::Idle - }, - None, - ) - .map_err(|_| internal_error())?; - - match result { - Ok(stats) => { - self.notify_compaction(session_id, &session, &connection, stats, 0)?; - Ok(PromptResponse::new(StopReason::EndTurn)) - } - Err(_error) if cancellation.is_cancelled() => { - Ok(PromptResponse::new(StopReason::Cancelled)) - } - Err(error) => Err(error), - } - } - - async fn btw_session( - &self, - session_id: &str, - session: AcpSession, - question: String, - connection: ConnectionTo, - cancellation: CancellationToken, - ) -> Result { - let history = self - .session_manager - .lock() - .map_err(|_| internal_error())? - .get_session_ref(session_id) - .map(|stored| stored.messages.clone()) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - let answer = { - let result = cancellable( - &cancellation, - crate::llm::client::generate_btw_answer( - session.provider, - session.model, - question, - history, - ), - ) - .await; - let Some(result) = result else { - return Ok(PromptResponse::new(StopReason::Cancelled)); - }; - result.map_err(|error| internal_error_with(&error.to_string()))? - }; - // Serialize notification with close/delete, which cancel under this same lock. - let sessions = self.sessions.lock().await; - if cancellation.is_cancelled() || !sessions.attachments.contains_key(session_id) { - return Ok(PromptResponse::new(StopReason::Cancelled)); - } - send_text(&connection, session_id, &cuid2::create_id(), answer, false)?; - Ok(PromptResponse::new(StopReason::EndTurn)) - } - - async fn run_compaction( - &self, - session_id: &str, - session: &AcpSession, - cancellation: CancellationToken, - minimum_tokens: usize, - ) -> Result { - let messages = { - let manager = self.session_manager.lock().map_err(|_| internal_error())?; - manager - .get_session_ref(session_id) - .map(|stored| stored.messages.clone()) - .ok_or_else(|| Error::invalid_params().data("unknown session"))? - }; - let selection = crate::session::compaction::select_messages_for_compaction_with_min( - &messages, - crate::session::compaction::DEFAULT_TAIL_TURNS, - minimum_tokens, - ) - .ok_or_else(|| Error::invalid_params().data("Nothing to compact"))?; - let before_tokens = crate::session::compaction::total_context_tokens(&messages); - let before_messages = - crate::session::compaction::filter_messages_for_context(&messages).len(); - let prompt = crate::session::compaction::build_prompt(&selection.messages_to_summarize); - let summary = cancellable( - &cancellation, - crate::llm::client::summarize_for_compaction( - session.provider.clone(), - session.model.clone(), - compaction_reasoning(session), - prompt, - cancellation.clone(), - ), - ) - .await - .ok_or_else(|| internal_error_with("Compaction cancelled by user"))? - .map_err(|error| internal_error_with(&error.to_string()))?; - if cancellation.is_cancelled() { - return Err(internal_error_with("Compaction cancelled by user")); - } - let (compacted, stats) = compacted_messages( - &messages, - &selection, - &summary, - session, - before_tokens, - before_messages, - )?; - let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; - manager - .replace_session_messages(session_id, compacted) - .map_err(|_| internal_error())?; - Ok(stats) - } - - async fn maybe_auto_compact_session( - &self, - session_id: &str, - session: &AcpSession, - connection: &ConnectionTo, - cancellation: CancellationToken, - prompt_tokens: usize, - additional_tokens: usize, - ) -> Result { - let messages = self - .session_manager - .lock() - .map_err(|_| internal_error())? - .get_session_ref(session_id) - .map(|stored| stored.messages.clone()) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - let used_tokens = - acp_context_tokens(&messages, prompt_tokens).saturating_add(additional_tokens); - if !crate::session::compaction::should_auto_compact( - &session.config.merged_config.compaction, - used_tokens, - session.context_window, - model_output_limit(&session.config, &session.provider, &session.model), - ) { - return Ok(false); - } - if crate::session::compaction::select_messages_for_compaction( - &messages, - crate::session::compaction::DEFAULT_TAIL_TURNS, - ) - .is_none() - { - return Ok(false); - } - self.session_manager - .lock() - .map_err(|_| internal_error())? - .set_session_status( - session_id, - crate::session::types::SessionStatus::Streaming, - None, - ) - .map_err(|_| internal_error())?; - let result = self - .run_compaction( - session_id, - session, - cancellation.clone(), - crate::session::compaction::MIN_COMPACTABLE_TOKENS, - ) - .await; - self.session_manager - .lock() - .map_err(|_| internal_error())? - .set_session_status( - session_id, - if cancellation.is_cancelled() { - crate::session::types::SessionStatus::Interrupted - } else { - crate::session::types::SessionStatus::Idle - }, - None, - ) - .map_err(|_| internal_error())?; - match result { - Ok(stats) => { - self.notify_compaction(session_id, session, connection, stats, prompt_tokens)?; - Ok(true) - } - Err(error) => { - crate::emit_log!("ACP auto-compaction skipped after failure: {:?}", error); - Ok(false) - } - } - } - - fn notify_compaction( - &self, - session_id: &str, - session: &AcpSession, - connection: &ConnectionTo, - stats: crate::session::types::CompactionStats, - prompt_tokens: usize, - ) -> Result<(), Error> { - let (used_tokens, cost) = { - let manager = self.session_manager.lock().map_err(|_| internal_error())?; - let stored = manager - .get_session_ref(session_id) - .ok_or_else(|| Error::invalid_params().data("unknown session"))?; - ( - acp_context_tokens(&stored.messages, prompt_tokens), - acp_session_cost(&stored.messages), - ) - }; - send_text( - connection, - session_id, - &cuid2::create_id(), - format!( - "Context compacted ({})", - crate::session::compaction::format_compaction_stats(stats) - ), - false, - )?; - send_usage( - connection, - session_id, - session, - used_tokens, - (cost > 0.0).then_some(cost), - None, - ) - } -} - -fn acp_session_cost(messages: &[crate::session::types::Message]) -> f64 { - messages - .iter() - .map(|message| message.cost.unwrap_or_else(|| message.usage_cost())) - .sum() -} - -fn acp_request_messages( - history: &[crate::session::types::Message], - system_message: crate::session::types::Message, -) -> Vec { - let mut messages = crate::session::compaction::filter_messages_for_context(history); - messages.insert(0, system_message); - messages -} - -fn acp_context_tokens(messages: &[crate::session::types::Message], prompt_tokens: usize) -> usize { - let estimated = - crate::session::compaction::total_context_tokens(messages).saturating_add(prompt_tokens); - // Retained tail usage still describes the old context, as does the summary's billed usage. - let start = messages - .iter() - .rposition(crate::session::compaction::is_compaction_marker) - .map(|index| index + 1) - .unwrap_or(0); - for (index, message) in messages.iter().enumerate().skip(start).rev() { - if message.role != crate::session::types::MessageRole::Assistant { - continue; - } - // Each tool step bills its prompt again; cumulative usage is not a context size. - let provider_context = message - .parts - .iter() - .enumerate() - .filter_map(|(part_index, part)| { - let usage = part.recorded_usage()?; - let input = usage - .input - .saturating_add(usage.cache_read) - .saturating_add(usage.cache_write); - if input == 0 { - return None; - } - let mut tokens = - usize::try_from(input.saturating_add(usage.output)).unwrap_or(usize::MAX); - if let Some(previous_prompt) = part - .data - .get("acp_prompt_tokens") - .and_then(serde_json::Value::as_u64) - { - tokens = tokens - .saturating_sub(usize::try_from(previous_prompt).unwrap_or(usize::MAX)) - .saturating_add(prompt_tokens); - } - Some(tokens.saturating_add( - crate::session::compaction::assistant_parts_context_tokens( - &message.parts[part_index + 1..], - ), - )) - }) - .max() - .filter(|tokens| *tokens > 0); - if let Some(tokens) = provider_context { - let following = - crate::session::compaction::total_context_tokens(&messages[index + 1..]); - return estimated.max(tokens.saturating_add(following)); - } - } - estimated -} - -fn compaction_reasoning(session: &AcpSession) -> Option { - use crate::model::reasoning::ReasoningEffort; - let capability = model_reasoning_capability( - &session.config, - &session.models, - &session.provider, - &session.model, - )?; - [ - ReasoningEffort::None, - ReasoningEffort::Minimal, - ReasoningEffort::Low, - ] - .into_iter() - .find(|effort| capability.values().contains(effort)) - .or(session.reasoning) - .filter(|effort| *effort != ReasoningEffort::None) -} - -fn compacted_messages( - messages: &[crate::session::types::Message], - selection: &crate::session::compaction::CompactionSelection, - summary: &crate::llm::client::CompactionSummary, - session: &AcpSession, - before_tokens: usize, - before_messages: usize, -) -> Result< - ( - Vec, - crate::session::types::CompactionStats, - ), - Error, -> { - let mut compacted = crate::session::compaction::apply_soft_compaction( - messages, - selection, - &summary.text, - Some(session.model.clone()), - Some(session.provider.clone()), - Some(session.agent.clone()), - crate::session::types::CompactionStats { - before_tokens, - after_tokens: 0, - before_messages, - after_messages: 0, - }, - ); - crate::session::compaction::attach_summary_usage( - &mut compacted, - crate::session::types::RecordedUsage { - input: summary.usage.input, - output: summary.usage.output, - cache_read: summary.usage.cache_read, - cache_write: summary.usage.cache_write, - cost: estimate_session_usage_cost(session, &summary.usage), - }, - ); - let after_tokens = crate::session::compaction::total_context_tokens(&compacted); - let after_messages = crate::session::compaction::filter_messages_for_context(&compacted).len(); - let stats = crate::session::types::CompactionStats { - before_tokens, - after_tokens, - before_messages, - after_messages, - }; - if after_tokens >= before_tokens { - return Err(Error::invalid_params().data(format!( - "Compaction did not reduce context ({})", - crate::session::compaction::format_compaction_stats(stats) - ))); - } - if let Some(marker) = compacted - .iter_mut() - .rev() - .find(|message| crate::session::compaction::is_compaction_marker(message)) - { - marker.compaction_stats = Some(stats); - } - Ok((compacted, stats)) -} - -fn acp_stop_reason(reason: Option) -> StopReason { - match reason { - Some(crate::llm::TurnStopReason::MaxTokens) => StopReason::MaxTokens, - Some(crate::llm::TurnStopReason::MaxTurnRequests) => StopReason::MaxTurnRequests, - Some(crate::llm::TurnStopReason::Refusal) => StopReason::Refusal, - None => StopReason::EndTurn, - } -} - -fn resolve_model(config: &LoadedConfig) -> (String, String) { - if let Ok(Some((provider, model))) = - crate::persistence::PrefsDAO::new().and_then(|prefs| prefs.get_active_model()) - { - return (provider, model); - } - config - .merged_config - .model - .as_deref() - .map(crate::app::parse_model_ref) - .unwrap_or_else(|| ("opencode".to_string(), "big-pickle".to_string())) -} - -fn configured_tool_permissions(cwd: &Path, config: &LoadedConfig) -> crate::tools::ToolPermissions { - let mut policies = crate::tools::AgentToolPolicies::default(); - for (mode, tools) in config.merged_config.agent_registry.tool_policy_map() { - policies = policies.with_custom_tools(mode, tools); - } - crate::tools::ToolPermissions::new(cwd) - .with_agent_policies(policies) - .with_permission_rules(config.merged_config.permission_rules.clone()) - .with_agent_permission_rules(config.merged_config.agent_registry.permission_rules_map()) -} - -fn tool_permissions(session: &AcpSession) -> crate::tools::ToolPermissions { - session.tool_permissions.clone() -} - -fn mode_description(agent: &crate::agent::definition::AgentDefinition) -> &str { - if agent.name == "compaction" && agent.description.trim().is_empty() { - "Compaction agent. Summarizes conversation context for continued work." - } else { - &agent.description - } -} - -fn session_modes(session: &AcpSession) -> SessionModeState { - let mut modes = session - .config - .merged_config - .agent_registry - .visible_primary_agents() - .into_iter() - .map(|agent| { - SessionMode::new(agent.name.clone(), agent.name.clone()) - .description(mode_description(agent)) - }) - .collect::>(); - modes.sort_by(|left, right| left.name.cmp(&right.name)); - SessionModeState::new(session.agent.clone(), modes) -} - -fn session_config_options(session: &AcpSession) -> Vec { - let mut mode_options = session - .config - .merged_config - .agent_registry - .visible_primary_agents() - .into_iter() - .map(|agent| { - SessionConfigSelectOption::new(agent.name.clone(), agent.name.clone()) - .description(mode_description(agent)) - }) - .collect::>(); - mode_options.sort_by(|left, right| left.name.cmp(&right.name)); - - let mut options = vec![ - SessionConfigOption::select("mode", "Mode", session.agent.clone(), mode_options) - .category(SessionConfigOptionCategory::Mode), - model_config_option(&session.models, &session.provider, &session.model), - ]; - options.push(reasoning_config_option(session)); - options -} - -fn reasoning_config_option(session: &AcpSession) -> SessionConfigOption { - let options = [ - crate::model::reasoning::ReasoningEffort::None, - crate::model::reasoning::ReasoningEffort::Low, - crate::model::reasoning::ReasoningEffort::Medium, - crate::model::reasoning::ReasoningEffort::High, - crate::model::reasoning::ReasoningEffort::XHigh, - crate::model::reasoning::ReasoningEffort::Max, - ] - .iter() - .map(|effort| SessionConfigSelectOption::new(effort.as_str(), reasoning_effort_label(*effort))) - .collect::>(); - let current = session.reasoning_selection.as_str(); - SessionConfigOption::select("effort", "Effort", current, options) - .category(SessionConfigOptionCategory::ThoughtLevel) -} - -fn model_config_option( - models: &[crate::model::types::Model], - provider: &str, - model_id: &str, -) -> SessionConfigOption { - let mut model_groups: Vec = Vec::new(); - for model in models { - let option = SessionConfigSelectOption::new(model_value(model), model.name.clone()) - .description(model.id.clone()); - if model_groups - .last() - .is_some_and(|group| group.group.to_string() == model.provider_id) - { - model_groups - .last_mut() - .expect("model group exists") - .options - .push(option); - } else { - model_groups.push(SessionConfigSelectGroup::new( - model.provider_id.clone(), - model.provider_name.clone(), - vec![option], - )); - } - } - - SessionConfigOption::select( - "model", - "Model", - format!("{provider}/{model_id}"), - model_groups, - ) - .category(SessionConfigOptionCategory::Model) -} - -fn model_value(model: &crate::model::types::Model) -> String { - crate::model::catalog::model_ref(model) -} - -fn find_selectable_model<'a>( - models: &'a [crate::model::types::Model], - model_ref: &str, -) -> Result<&'a crate::model::types::Model, Error> { - models - .iter() - .find(|model| model_value(model) == model_ref) - .ok_or_else(|| Error::invalid_params().data("unknown or unavailable ACP model")) -} - -fn model_reasoning( - config: &LoadedConfig, - models: &[crate::model::types::Model], - provider: &str, - model_id: &str, -) -> Option { - model_reasoning_capability(config, models, provider, model_id) - .and_then(|capability| capability.resolve(None)) -} - -fn model_reasoning_capability( - config: &LoadedConfig, - models: &[crate::model::types::Model], - provider: &str, - model_id: &str, -) -> Option { - if let Some(capability) = models - .iter() - .find(|model| model.provider_id == provider && model.id == model_id) - .and_then(|model| { - crate::model::reasoning::capability_from_options(&model.reasoning_options) - }) - { - return Some(capability); - } - crate::model::discovery::Discovery::new_with_custom(Some( - config.merged_config.custom_providers.clone(), - )) - .ok() - .and_then(|discovery| discovery.get_model_reasoning_capability(provider, model_id)) - .filter(|capability| !capability.values().is_empty()) -} - -fn estimate_session_usage_cost( - session: &AcpSession, - usage: &crate::aisdk::chunk::TokenUsage, -) -> f64 { - crate::model::discovery::Discovery::new_with_custom(Some( - session.config.merged_config.custom_providers.clone(), - )) - .ok() - .map(|discovery| { - discovery.estimate_usage_cost( - &session.provider, - &session.model, - usage.input, - usage.output, - usage.cache_read, - usage.cache_write, - ) - }) - .unwrap_or(0.0) -} - -fn model_context_window( - config: &LoadedConfig, - models: &[crate::model::types::Model], - provider: &str, - model: &str, -) -> Option { - if let Some(context_window) = models - .iter() - .find(|candidate| candidate.provider_id == provider && candidate.id == model) - .and_then(|model| model.context_window) - { - return Some(context_window); - } - - let discovery = crate::model::discovery::Discovery::new_with_custom(Some( - config.merged_config.custom_providers.clone(), - )) - .ok(); - discovery - .as_ref() - .and_then(|discovery| discovery.get_model_limit(provider, model)) - .or_else(|| { - config - .merged_config - .custom_providers - .get(provider) - .and_then(|provider| provider.models.get(model)) - .and_then(|model| model.context_window) - }) -} - -fn workspace_path(path: &Path) -> Result { - if !path.is_absolute() { - return Err(Error::invalid_params().data("cwd must be an absolute path")); - } - let path = path - .canonicalize() - .map_err(|_| Error::resource_not_found(Some(path.to_string_lossy().to_string())))?; - if !path.is_dir() { - return Err(Error::invalid_params().data("cwd must be a directory")); - } - Ok(path) -} - -fn prompt_content( - parts: Vec, - text_override: Option<&str>, - supports_images: bool, - supports_audio: bool, - session_id: &str, - session: &AcpSession, -) -> Result<(String, Vec, Vec), Error> { - let mut text = String::new(); - let mut text_override = text_override; - let mut local_image_paths = Vec::new(); - let mut local_audio_paths = Vec::new(); - for part in parts { - match part { - ContentBlock::Text(content) => { - text.push_str(text_override.take().unwrap_or(&content.text)); - } - ContentBlock::ResourceLink(link) => { - text.push_str(&format!("[{}]", link.uri)); - } - ContentBlock::Resource(resource) => match resource.resource { - EmbeddedResourceResource::TextResourceContents(resource) => { - text.push_str(&format!("[{}]\n{}", resource.uri, resource.text)); - } - EmbeddedResourceResource::BlobResourceContents(resource) => { - text.push_str(&format!("[{}]", resource.uri)); - } - _ => {} - }, - ContentBlock::Image(image) => { - if !supports_images { - return Err(Error::invalid_params().data(format!( - "model {}/{} does not support image input", - session.provider, session.model - ))); - } - match write_prompt_image(session_id, &image) { - Ok(path) => local_image_paths.push(path), - Err(error) => { - for path in &local_image_paths { - crate::persistence::attachments::remove_file(Path::new(path)); - } - return Err(error); - } - } - } - ContentBlock::Audio(audio) => { - if !supports_audio { - for path in local_image_paths.iter().chain(local_audio_paths.iter()) { - crate::persistence::attachments::remove_file(Path::new(path)); - } - return Err(Error::invalid_params().data(format!( - "model {}/{} does not support audio input", - session.provider, session.model - ))); - } - match write_prompt_audio(session_id, &audio) { - Ok(path) => local_audio_paths.push(path), - Err(error) => { - for path in local_image_paths.iter().chain(local_audio_paths.iter()) { - crate::persistence::attachments::remove_file(Path::new(path)); - } - return Err(error); - } - } - } - _ => {} - } - } - if text.is_empty() { - if !local_image_paths.is_empty() { - text.push_str("[Image attached]"); - } else if !local_audio_paths.is_empty() { - text.push_str("[Audio attached]"); - } - } - Ok((text, local_image_paths, local_audio_paths)) -} - -fn prompt_text(parts: Vec) -> Result { - let mut text = String::new(); - for part in parts { - match part { - ContentBlock::Text(content) => text.push_str(&content.text), - ContentBlock::ResourceLink(link) => text.push_str(&format!("[{}]", link.uri)), - ContentBlock::Resource(resource) => match resource.resource { - EmbeddedResourceResource::TextResourceContents(resource) => { - text.push_str(&format!("[{}]\n{}", resource.uri, resource.text)); - } - EmbeddedResourceResource::BlobResourceContents(resource) => { - text.push_str(&format!("[{}]", resource.uri)); - } - _ => {} - }, - ContentBlock::Image(_) | ContentBlock::Audio(_) => { - return Err( - Error::invalid_params().data("binary ACP prompt content is not supported here") - ); - } - _ => {} - } - } - Ok(text) -} - -fn write_prompt_image( - session_id: &str, - image: &agent_client_protocol::schema::v1::ImageContent, -) -> Result { - const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024; - - let (mime_type, encoded_data) = prompt_image_payload(image)?; - let extension = match mime_type { - "image/png" => "png", - "image/jpeg" => "jpg", - "image/gif" => "gif", - "image/webp" => "webp", - mime_type => { - return Err( - Error::invalid_params().data(format!("unsupported image MIME type: {mime_type}")) - ); - } - }; - let data = base64::engine::general_purpose::STANDARD - .decode(encoded_data) - .map_err(|error| Error::invalid_params().data(format!("invalid image data: {error}")))?; - if data.len() > MAX_IMAGE_BYTES { - return Err(Error::invalid_params().data("image exceeds the 20 MiB size limit")); - } - - let path = crate::persistence::attachments::write(session_id, extension, &data) - .map_err(|_| internal_error())?; - Ok(path.to_string_lossy().into_owned()) -} - -fn prompt_image_payload( - image: &agent_client_protocol::schema::v1::ImageContent, -) -> Result<(&str, &str), Error> { - let data = image.data.trim(); - if data.is_empty() { - return Err(Error::invalid_params().data("image data is empty")); - } - - let Some(data_uri) = data.strip_prefix("data:") else { - return Ok((image.mime_type.as_str(), data)); - }; - let (metadata, payload) = data_uri - .split_once(',') - .ok_or_else(|| Error::invalid_params().data("invalid image data URI"))?; - let mut metadata = metadata.split(';'); - let mime_type = metadata - .next() - .filter(|mime_type| !mime_type.is_empty()) - .ok_or_else(|| Error::invalid_params().data("image data URI is missing a MIME type"))?; - if !metadata.any(|value| value.eq_ignore_ascii_case("base64")) { - return Err(Error::invalid_params().data("image data URI must be base64 encoded")); - } - let payload = payload.trim(); - if payload.is_empty() { - return Err(Error::invalid_params().data("image data is empty")); - } - - Ok((mime_type, payload)) -} - -fn send_text( - connection: &ConnectionTo, - session_id: &str, - message_id: &str, - text: String, - thought: bool, -) -> Result<(), Error> { - let update = if thought { - SessionUpdate::AgentThoughtChunk(ContentChunk::new(text.into()).message_id(message_id)) - } else { - SessionUpdate::AgentMessageChunk(ContentChunk::new(text.into()).message_id(message_id)) - }; - connection - .send_notification(SessionNotification::new(session_id.to_string(), update)) - .map_err(|_| internal_error()) -} - -fn tool_result_text(payload: &serde_json::Value) -> String { - payload - .get("output") - .or_else(|| payload.get("output_preview")) - .or_else(|| payload.get("error")) - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string() -} - -fn absolute_tool_path(path: &str, cwd: &Path) -> PathBuf { - let path = PathBuf::from(path); - if path.is_absolute() { - path - } else { - cwd.join(path) - } -} - -fn tool_locations(tool_name: &str, input: &serde_json::Value, cwd: &Path) -> Vec { - let paths = if tool_name == "write_files" { - input - .get("files") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(|file| file.get("file_path").and_then(serde_json::Value::as_str)) - .map(str::to_string) - .collect() - } else if tool_name == "apply_patch" { - crate::tools::patch::patch_paths_from_params(input) - } else { - input - .get("file_path") - .or_else(|| input.get("filePath")) - .or_else(|| input.get("filepath")) - .or_else(|| input.get("path")) - .and_then(serde_json::Value::as_str) - .map(|path| vec![path.to_string()]) - .unwrap_or_default() - }; - - paths - .into_iter() - .map(|path| ToolCallLocation::new(absolute_tool_path(&path, cwd))) - .collect() -} - -fn tool_result_locations(payload: &serde_json::Value, cwd: &Path) -> Vec { - let Some(metadata) = payload.get("metadata") else { - return Vec::new(); - }; - let line = metadata - .get("line_number") - .and_then(serde_json::Value::as_u64) - .and_then(|line| u32::try_from(line).ok()); - - if let Some(changes) = metadata - .get("changes") - .and_then(serde_json::Value::as_array) - { - return changes - .iter() - .filter_map(|change| change.get("path").and_then(serde_json::Value::as_str)) - .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd))) - .collect(); - } - - metadata - .get("path") - .and_then(serde_json::Value::as_str) - .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd)).line(line)) - .into_iter() - .collect() -} - -fn tool_result_content(payload: &serde_json::Value, cwd: &Path) -> Vec { - let mut content = Vec::new(); - let text = tool_result_text(payload); - if !text.is_empty() { - content.push(ToolCallContent::from(text)); - } - - if let Some(metadata) = payload.get("metadata") { - if let Some(changes) = metadata - .get("changes") - .and_then(serde_json::Value::as_array) - { - content.extend(changes.iter().filter_map(|change| tool_diff(change, cwd))); - } else if let Some(diff) = tool_diff(metadata, cwd) { - content.push(diff); - } - content.extend(mcp_native_tool_content(metadata)); - } - - if !metadata_has_mcp_images(payload.get("metadata")) { - if let Some(images) = payload.get("images").and_then(serde_json::Value::as_array) { - content.extend(images.iter().filter_map(tool_result_image)); - } - } - - content -} - -fn tool_diff(change: &serde_json::Value, cwd: &Path) -> Option { - let path = change.get("path")?.as_str()?; - let new_text = change.get("new_text")?.as_str()?; - let old_text = change.get("old_text").and_then(serde_json::Value::as_str); - Some( - agent_client_protocol::schema::v1::Diff::new( - absolute_tool_path(path, cwd), - new_text.to_string(), - ) - .old_text(old_text.map(str::to_string)) - .into(), - ) -} - -fn tool_result_image(image: &serde_json::Value) -> Option { - let data = image.get("data_url")?.as_str()?; - let media_type = image.get("media_type")?.as_str()?; - let encoded = data - .strip_prefix("data:") - .and_then(|value| value.split_once(',')) - .map(|(_, encoded)| encoded) - .unwrap_or(data); - Some(ToolCallContent::from(ContentBlock::Image( - agent_client_protocol::schema::v1::ImageContent::new(encoded, media_type), - ))) -} - -fn send_tool_call( - connection: &ConnectionTo, - session_id: &str, - tool_call: crate::llm::ToolCall, - cwd: &Path, -) -> Result<(), Error> { - let raw_input = serde_json::from_str(&tool_call.function.arguments) - .unwrap_or_else(|_| serde_json::json!({ "arguments": tool_call.function.arguments })); - let title = tool_title(&tool_call.function.name, &raw_input); - let update = SessionUpdate::ToolCall( - ToolCall::new(tool_call.id, title) - .kind(tool_kind(&tool_call.function.name)) - .status(ToolCallStatus::Pending) - .locations(tool_locations(&tool_call.function.name, &raw_input, cwd)) - .raw_input(raw_input), - ); - connection - .send_notification(SessionNotification::new(session_id.to_string(), update)) - .map_err(|_| internal_error()) -} - -async fn request_permission( - connection: &ConnectionTo, - session_id: &str, - prompt: &crate::tools::PermissionPrompt, -) -> crate::tools::PermissionResponse { - let tool_call_id = permission_tool_call_id(prompt.tool_call_id.as_deref()); - let input = serde_json::json!({ - "tool": prompt.tool_id, - "permission": prompt.permission, - "patterns": prompt.patterns, - "target": prompt.target, - "command": prompt.command, - "workdir": prompt.workdir, - "reason": prompt.reason, - "input": prompt.raw_input, - }); - let cwd = PathBuf::from(&prompt.workspace); - let content = permission_tool_content(prompt, &cwd); - let tool_call = ToolCallUpdate::new(tool_call_id, { - let mut fields = ToolCallUpdateFields::new() - .title(permission_title(prompt)) - .kind(tool_kind(&prompt.tool_id)) - .status(ToolCallStatus::Pending) - .locations(tool_locations(&prompt.tool_id, &prompt.raw_input, &cwd)) - .raw_input(input); - if !content.is_empty() { - fields = fields.content(content); - } - fields - }); - let request = - RequestPermissionRequest::new(session_id.to_string(), tool_call, permission_options()); - let Ok(response) = connection.send_request(request).block_task().await else { - return crate::tools::PermissionResponse::Deny; - }; - match response.outcome { - RequestPermissionOutcome::Selected(selection) - if selection.option_id.to_string() == "once" => - { - crate::tools::PermissionResponse::AllowOnce - } - RequestPermissionOutcome::Selected(selection) - if selection.option_id.to_string() == "always" => - { - crate::tools::PermissionResponse::AllowAlways - } - RequestPermissionOutcome::Selected(_) | RequestPermissionOutcome::Cancelled => { - crate::tools::PermissionResponse::Deny - } - _ => crate::tools::PermissionResponse::Deny, - } -} - -async fn request_questions( - connection: &ConnectionTo, - session_id: &str, - tool_call_id: Option<&str>, - questions: &serde_json::Value, - cancellation: &CancellationToken, -) -> serde_json::Value { - let form = acp_question_form(session_id, tool_call_id, questions); - let request = connection.send_request(form.request).block_task(); - tokio::pin!(request); - let response = tokio::select! { - _ = cancellation.cancelled() => return skipped_question_answers(questions), - response = &mut request => response, - }; - let Ok(response) = response else { - return skipped_question_answers(questions); - }; - acp_question_answers(&form.fields, response.action) -} - -async fn bridge_terminal_session( - connection: &ConnectionTo, - session_id: &str, - session_cwd: &Path, - request: crate::tools::TerminalSessionRequest, - cancellation: &CancellationToken, -) { - // Connection-owned task retains the create response even when the prompt stops. - // Do not tokio::spawn SDK block_task futures: they need the connection task context. - let worker_connection = connection.clone(); - let worker_session_id = session_id.to_string(); - let worker_cwd = session_cwd.to_path_buf(); - let worker_cancellation = cancellation.clone(); - let control_tx = request.control_tx.clone(); - let (done_tx, done_rx) = tokio::sync::oneshot::channel(); - if connection - .spawn(async move { - terminal_session_worker( - &worker_connection, - &worker_session_id, - &worker_cwd, - request, - &worker_cancellation, - ) - .await; - let _ = done_tx.send(()); - Ok(()) - }) - .is_err() - { - let _ = control_tx.send(crate::tools::TerminalSessionControl::Stop); - return; - } - if cancellable(cancellation, done_rx).await.is_none() { - let _ = control_tx.send(crate::tools::TerminalSessionControl::Stop); - } -} - -async fn cleanup_terminal( - connection: &ConnectionTo, - session_id: &str, - terminal_id: agent_client_protocol::schema::v1::TerminalId, -) { - // Issue both requests even if a client never acknowledges kill. - let kill = connection - .send_request(KillTerminalRequest::new( - session_id.to_string(), - terminal_id.clone(), - )) - .block_task(); - let release = connection - .send_request(ReleaseTerminalRequest::new( - session_id.to_string(), - terminal_id, - )) - .block_task(); - let _ = tokio::join!(kill, release); -} - -async fn terminal_session_worker( - connection: &ConnectionTo, - session_id: &str, - session_cwd: &Path, - request: crate::tools::TerminalSessionRequest, - cancellation: &CancellationToken, -) { - let start = request.start; - let control_tx = request.control_tx; - let cwd = start - .workdir - .as_deref() - .map(PathBuf::from) - .map(|path| absolute_tool_path(&path.to_string_lossy(), session_cwd)) - .unwrap_or_else(|| session_cwd.to_path_buf()); - let create = CreateTerminalRequest::new(session_id.to_string(), "bash") - .args(vec!["-c".to_string(), start.command.clone()]) - .cwd(cwd) - .output_byte_limit(crate::tools::terminal_session::MAX_TRANSCRIPT_BYTES as u64); - let terminal_id = match connection.send_request(create).block_task().await { - Ok(response) => response.terminal_id, - Err(error) => { - let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( - format!("ACP client could not create terminal: {error}"), - )); - return; - } - }; - - if cancellation.is_cancelled() { - cleanup_terminal(connection, session_id, terminal_id).await; - return; - } - - let terminal_update = SessionUpdate::ToolCallUpdate(ToolCallUpdate::new( - start.tool_call_id.clone(), - ToolCallUpdateFields::new().content(vec![ToolCallContent::Terminal(Terminal::new( - terminal_id.clone(), - ))]), - )); - if connection - .send_notification(SessionNotification::new( - session_id.to_string(), - terminal_update, - )) - .is_err() - { - cleanup_terminal(connection, session_id, terminal_id).await; - let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( - "ACP client could not embed terminal".to_string(), - )); - return; - } - - let wait = connection - .send_request(WaitForTerminalExitRequest::new( - session_id.to_string(), - terminal_id.clone(), - )) - .block_task(); - tokio::pin!(wait); - let (stopped_by_user, exit_code, wait_error) = tokio::select! { - _ = cancellation.cancelled() => { - cleanup_terminal(connection, session_id, terminal_id).await; - return; - } - response = &mut wait => match response { - Ok(response) => ( - false, - response.exit_status.exit_code.and_then(|code| i32::try_from(code).ok()), - None, - ), - Err(error) => (false, None, Some(error.to_string())), - } - }; - - let output = cancellable( - cancellation, - connection - .send_request(TerminalOutputRequest::new( - session_id.to_string(), - terminal_id.clone(), - )) - .block_task(), - ) - .await; - let Some(output) = output else { - cleanup_terminal(connection, session_id, terminal_id).await; - return; - }; - let _ = connection - .send_request(ReleaseTerminalRequest::new( - session_id.to_string(), - terminal_id, - )) - .block_task() - .await; - - match (wait_error, output) { - (None, Ok(output)) => { - let result = crate::tools::terminal_session::external_terminal_result( - &start, - &output.output, - output.truncated, - exit_code, - stopped_by_user, - ); - let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalResult(result)); - } - (Some(error), _) => { - let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( - format!("ACP terminal wait failed: {error}"), - )); - } - (None, Err(error)) => { - let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( - format!("ACP terminal output failed: {error}"), - )); - } - } -} - -fn permission_title(prompt: &crate::tools::PermissionPrompt) -> String { - prompt - .command - .as_deref() - .or(prompt.target.as_deref()) - .unwrap_or(&prompt.reason) - .to_string() -} - -fn send_tool_result( - connection: &ConnectionTo, - session_id: &str, - result: crate::llm::ToolCallResult, - cwd: &Path, -) -> Result<(), Error> { - let payload = serde_json::from_str::(&result.content).unwrap_or_else( - |_| serde_json::json!({ "status": "error", "output_preview": result.content }), - ); - let status = match payload.get("status").and_then(serde_json::Value::as_str) { - Some("ok") => ToolCallStatus::Completed, - _ => ToolCallStatus::Failed, - }; - let content = tool_result_content(&payload, cwd); - let locations = tool_result_locations(&payload, cwd); - let mut fields = ToolCallUpdateFields::new() - .status(status) - .content((!content.is_empty()).then_some(content)) - .raw_output(payload); - if !locations.is_empty() { - fields = fields.locations(locations); - } - let update = SessionUpdate::ToolCallUpdate(ToolCallUpdate::new(result.tool_call_id, fields)); - connection - .send_notification(SessionNotification::new(session_id.to_string(), update)) - .map_err(|_| internal_error()) -} - -fn send_usage( - connection: &ConnectionTo, - session_id: &str, - session: &AcpSession, - used: usize, - cost: Option, - usage: Option, -) -> Result<(), Error> { - let update = SessionUpdate::UsageUpdate(usage_update(session, used, cost, usage)); - connection - .send_notification(SessionNotification::new(session_id.to_string(), update)) - .map_err(|_| internal_error()) -} - -fn send_live_usage( - connection: &ConnectionTo, - session_id: &str, - session: &AcpSession, - base_context_tokens: usize, - base_cost: f64, - assistant: &crate::session::types::Message, - live_usage: &AcpLiveUsage, -) -> Result<(), Error> { - send_usage( - connection, - session_id, - session, - live_usage.context_tokens(base_context_tokens, assistant), - cumulative_cost(base_cost, live_usage.cost), - (!live_usage.cumulative.is_empty()).then_some(live_usage.cumulative), - ) -} - -#[derive(Debug, Default)] -struct AcpLiveUsage { - cumulative: crate::aisdk::chunk::TokenUsage, - provider_context_floor: usize, - cost: f64, -} - -impl AcpLiveUsage { - fn observe_provider_usage( - &mut self, - assistant: &crate::session::types::Message, - usage: crate::aisdk::chunk::TokenUsage, - ) { - self.cumulative = self.cumulative.saturating_add(usage); - let provider_input = usage - .input - .saturating_add(usage.cache_read) - .saturating_add(usage.cache_write); - let observed_assistant = live_assistant_context_tokens(assistant); - self.provider_context_floor = self - .provider_context_floor - .max(usize::try_from(provider_input).unwrap_or(usize::MAX)) - .max(observed_assistant); - } - - fn context_tokens( - &self, - base_context_tokens: usize, - assistant: &crate::session::types::Message, - ) -> usize { - base_context_tokens - .saturating_add(live_assistant_context_tokens(assistant)) - .max(self.provider_context_floor) - } -} - -fn live_assistant_context_tokens(assistant: &crate::session::types::Message) -> usize { - let persisted = crate::session::compaction::message_context_tokens(assistant); - let reasoning = assistant - .reasoning - .as_deref() - .map(estimate_acp_tokens) - .unwrap_or(0); - persisted.saturating_add(reasoning) -} - -fn estimate_acp_tokens(content: &str) -> usize { - content.chars().count().saturating_add(3) / 4 -} - -fn cumulative_cost(base_cost: f64, turn_cost: f64) -> Option { - let cost = base_cost + turn_cost; - (cost > 0.0).then_some(cost) -} - -fn record_tool_call( - assistant: &mut crate::session::types::Message, - tool_call: &crate::llm::ToolCall, -) { - let args = serde_json::from_str(&tool_call.function.arguments) - .unwrap_or_else(|_| serde_json::Value::String(tool_call.function.arguments.clone())); - let provider_executed = matches!( - tool_call.function.name.as_str(), - "x_search" | "web_search" | "file_search" - ); - - if let Some(part) = assistant - .parts - .iter_mut() - .find(|part| part.part_type == "tool_call" && part.tool_id() == Some(tool_call.id.as_str())) - { - if let Some(obj) = part.data.as_object_mut() { - obj.insert( - "name".to_string(), - serde_json::Value::String(tool_call.function.name.clone()), - ); - obj.insert("args".to_string(), args); - if provider_executed { - obj.insert( - "provider_executed".to_string(), - serde_json::Value::Bool(true), - ); - } - } - return; - } - - assistant.add_tool_call_part(tool_call.id.clone(), tool_call.function.name.clone(), args); - if provider_executed { - if let Some(obj) = assistant - .parts - .last_mut() - .and_then(|part| part.data.as_object_mut()) - { - obj.insert( - "provider_executed".to_string(), - serde_json::Value::Bool(true), - ); - } - } -} - -fn record_tool_result( - assistant: &mut crate::session::types::Message, - result: &crate::llm::ToolCallResult, -) { - let mut data = assistant - .tool_call_part_data(&result.tool_call_id) - .cloned() - .unwrap_or_else(|| serde_json::json!({})); - data["id"] = serde_json::Value::String(result.tool_call_id.clone()); - data["name"] = serde_json::Value::String(result.name.clone()); - if matches!( - result.name.as_str(), - "x_search" | "web_search" | "file_search" - ) { - data["provider_executed"] = serde_json::Value::Bool(true); - } - - match serde_json::from_str::(&result.content) { - Ok(serde_json::Value::Object(payload)) => { - if let Some(data) = data.as_object_mut() { - for (key, value) in payload { - data.insert(key, value); - } - } - } - _ => { - data["status"] = serde_json::Value::String("ok".to_string()); - data["output_preview"] = serde_json::Value::String(result.content.clone()); - } - } - assistant.add_or_update_tool_result_part(data); -} - -fn usage_update( - session: &AcpSession, - used: usize, - cost: Option, - usage: Option, -) -> UsageUpdate { - let size = session - .context_window - .map(u64::from) - .unwrap_or_else(|| (used as u64).max(1)); - let mut meta = serde_json::Map::new(); - meta.insert( - "crabcode".to_string(), - serde_json::json!({ - "contextWindowKnown": session.context_window.is_some(), - "usage": usage.map(|usage| serde_json::json!({ - "inputTokens": usage.input, - "outputTokens": usage.output, - "cacheReadTokens": usage.cache_read, - "cacheWriteTokens": usage.cache_write, - })), - }), - ); - UsageUpdate::new(used as u64, size) - .cost(cost.map(|amount| AcpCost::new(amount, "USD"))) - .meta(meta) -} - -fn replay_messages( - connection: &ConnectionTo, - session_id: &str, - messages: &[crate::session::types::Message], - cwd: &Path, -) -> Result<(), Error> { - for message in messages { - if crate::session::compaction::is_compaction_display_item(message) { - continue; - } - let message_id = message.id.clone(); - match message.role { - crate::session::types::MessageRole::User => { - if !message.content.is_empty() { - send_replay_text( - connection, - session_id, - &message_id, - &message.content, - true, - false, - )?; - } - } - crate::session::types::MessageRole::Assistant => { - for part in &message.parts { - match part.part_type.as_str() { - "text" => { - if let Some(text) = part.text_value() { - send_replay_text( - connection, - session_id, - &message_id, - text, - false, - false, - )?; - } - } - "reasoning" => { - if let Some(text) = part.text_value() { - send_replay_text( - connection, - session_id, - &message_id, - text, - false, - true, - )?; - } - } - "tool_call" => replay_tool_call(connection, session_id, part, cwd)?, - "tool_result" => replay_tool_result(connection, session_id, part, cwd)?, - _ => {} - } - } - } - crate::session::types::MessageRole::System - | crate::session::types::MessageRole::Tool => {} - } - } - Ok(()) -} - -fn send_replay_text( - connection: &ConnectionTo, - session_id: &str, - message_id: &str, - text: &str, - user: bool, - thought: bool, -) -> Result<(), Error> { - let chunk = ContentChunk::new(text.to_string().into()).message_id(message_id); - let update = if user { - SessionUpdate::UserMessageChunk(chunk) - } else if thought { - SessionUpdate::AgentThoughtChunk(chunk) - } else { - SessionUpdate::AgentMessageChunk(chunk) - }; - connection - .send_notification(SessionNotification::new(session_id.to_string(), update)) - .map_err(|_| internal_error()) -} - -fn replay_tool_call( - connection: &ConnectionTo, - session_id: &str, - part: &crate::session::types::MessagePart, - cwd: &Path, -) -> Result<(), Error> { - let Some(tool_call_id) = part.tool_id() else { - return Ok(()); - }; - let name = part.tool_name().unwrap_or("tool"); - let input = part - .data - .get("args") - .cloned() - .unwrap_or(serde_json::Value::Null); - let status = match part.tool_status() { - Some("completed") | Some("ok") => ToolCallStatus::Completed, - Some("error") | Some("failed") => ToolCallStatus::Failed, - Some("running") => ToolCallStatus::InProgress, - _ => ToolCallStatus::Pending, - }; - let update = SessionUpdate::ToolCall( - ToolCall::new(tool_call_id.to_string(), tool_title(name, &input)) - .kind(tool_kind(name)) - .status(status) - .locations(tool_locations(name, &input, cwd)) - .raw_input(input), - ); - connection - .send_notification(SessionNotification::new(session_id.to_string(), update)) - .map_err(|_| internal_error()) -} - -fn replay_tool_result( - connection: &ConnectionTo, - session_id: &str, - part: &crate::session::types::MessagePart, - cwd: &Path, -) -> Result<(), Error> { - let Some(tool_call_id) = part.tool_id() else { - return Ok(()); - }; - let content = part - .data - .get("content") - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string(); - send_tool_result( - connection, - session_id, - crate::llm::ToolCallResult { - tool_call_id: tool_call_id.to_string(), - role: "tool".to_string(), - name: part.tool_name().unwrap_or("tool").to_string(), - content, - }, - cwd, - ) -} - -fn tool_kind(tool_name: &str) -> ToolKind { - match tool_name { - "bash" | "bash_output" | "bash_kill" | "bash_restart" | "terminal_session" => { - ToolKind::Execute - } - "webfetch" => ToolKind::Fetch, - "grep" | "glob" | "context" => ToolKind::Search, - "read" | "view_image" => ToolKind::Read, - "edit" | "write" | "write_files" | "apply_patch" => ToolKind::Edit, - "task" => ToolKind::Think, - _ => ToolKind::Other, - } -} - -fn tool_title(tool_name: &str, input: &serde_json::Value) -> String { - match tool_name { - "bash" => input - .get("command") - .and_then(serde_json::Value::as_str) - .unwrap_or(tool_name) - .to_string(), - "read" | "edit" | "write" | "grep" | "glob" => input - .get("filePath") - .or_else(|| input.get("filepath")) - .or_else(|| input.get("path")) - .or_else(|| input.get("pattern")) - .and_then(serde_json::Value::as_str) - .unwrap_or(tool_name) - .to_string(), - _ => tool_name.to_string(), - } -} - -fn system_time_to_iso8601(value: std::time::SystemTime) -> String { - chrono::DateTime::::from(value).to_rfc3339() -} - -fn internal_error() -> Error { - Error::internal_error().data("Crabcode ACP operation failed") -} - -fn internal_error_with(error: &str) -> Error { - Error::internal_error().data(format!("Crabcode ACP operation failed: {error}")) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn model( - provider_id: &str, - provider_name: &str, - id: &str, - name: &str, - ) -> crate::model::types::Model { - crate::model::types::Model { - id: id.to_string(), - name: name.to_string(), - family: String::new(), - provider_id: provider_id.to_string(), - provider_name: provider_name.to_string(), - attachment: false, - structured_output: false, - free: false, - local: false, - reasoning_options: Vec::new(), - context_window: None, - } - } - - fn reasoning_model() -> crate::model::types::Model { - let mut model = model("openai", "OpenAI", "o3", "o3"); - model.reasoning_options = vec![crate::model::reasoning::ReasoningOption { - kind: "effort".to_string(), - values: vec!["low".to_string(), "medium".to_string(), "high".to_string()], - }]; - model - } - - fn test_session() -> AcpSession { - AcpSession { - cwd: PathBuf::from("/tmp"), - config: crate::config::configuration::LoadedConfig { - merged_config: crate::config::configuration::MergedConfig::default(), - raw_merged: serde_json::Value::Null, - diagnostics: Default::default(), - inventory: Default::default(), - project_root: PathBuf::from("/tmp"), - cwd: PathBuf::from("/tmp"), - xdg_config_home: PathBuf::from("/tmp"), - }, - tool_permissions: crate::tools::ToolPermissions::new("/tmp"), - skills: crate::skill::SkillStore::load(Path::new("/tmp"), Path::new("/tmp")), - models: vec![model("example", "Example", "chat", "Chat")], - provider: "example".to_string(), - model: "chat".to_string(), - agent: "Build".to_string(), - reasoning_selection: crate::model::reasoning::ReasoningEffort::High, - reasoning: None, - context_window: None, - cancellation: None, - } - } - - #[test] - fn maps_crabcode_tools_to_acp_kinds() { - assert_eq!(tool_kind("bash"), ToolKind::Execute); - assert_eq!(tool_kind("read"), ToolKind::Read); - assert_eq!(tool_kind("apply_patch"), ToolKind::Edit); - assert_eq!(tool_kind("unknown"), ToolKind::Other); - } - - #[test] - fn acp_usage_update_includes_cumulative_usd_cost() { - let usage = crate::aisdk::chunk::TokenUsage { - input: 800, - output: 200, - cache_read: 500, - cache_write: 100, - }; - let update = usage_update(&test_session(), 1_000, Some(0.125), Some(usage)); - assert_eq!(update.cost.as_ref().map(|cost| cost.amount), Some(0.125)); - assert_eq!( - update.cost.as_ref().map(|cost| cost.currency.as_str()), - Some("USD") - ); - assert_eq!(update.size, 1_000); - let meta = update.meta.expect("usage metadata"); - assert_eq!(meta["crabcode"]["contextWindowKnown"], false); - assert_eq!(meta["crabcode"]["usage"]["inputTokens"], 800); - assert_eq!(meta["crabcode"]["usage"]["cacheWriteTokens"], 100); - } - - #[test] - fn acp_live_usage_advances_with_streamed_tools_and_provider_context() { - let mut assistant = crate::session::types::Message::incomplete("hello"); - record_tool_call( - &mut assistant, - &crate::llm::ToolCall { - id: "call_1".to_string(), - call_type: "function".to_string(), - function: crate::llm::FunctionCall { - name: "read".to_string(), - arguments: serde_json::json!({ "filePath": "src/main.rs" }).to_string(), - }, - }, - ); - record_tool_result( - &mut assistant, - &crate::llm::ToolCallResult { - tool_call_id: "call_1".to_string(), - role: "tool".to_string(), - name: "read".to_string(), - content: serde_json::json!({ - "status": "ok", - "output_preview": "x".repeat(4_000), - }) - .to_string(), - }, - ); - - let mut live = AcpLiveUsage::default(); - let before_provider_usage = live.context_tokens(18_000, &assistant); - assert!(before_provider_usage > 18_000); - assert!(assistant - .tool_call_part_data("call_1") - .and_then(|part| part.get("args")) - .is_some()); - - live.observe_provider_usage( - &assistant, - crate::aisdk::chunk::TokenUsage { - input: 70_000, - output: 500, - cache_read: 20_000, - cache_write: 0, - }, - ); - assert_eq!(live.context_tokens(18_000, &assistant), 90_000); - assert_eq!(live.cumulative.input, 70_000); - assert_eq!(live.cumulative.cache_read, 20_000); - } - - #[test] - fn acp_live_usage_keeps_largest_provider_context_across_tool_steps() { - let assistant = crate::session::types::Message::incomplete(""); - let mut live = AcpLiveUsage::default(); - live.observe_provider_usage( - &assistant, - crate::aisdk::chunk::TokenUsage { - input: 60_000, - output: 200, - cache_read: 20_000, - cache_write: 0, - }, - ); - live.observe_provider_usage( - &assistant, - crate::aisdk::chunk::TokenUsage { - input: 10_000, - output: 100, - cache_read: 40_000, - cache_write: 0, - }, - ); - - assert_eq!(live.context_tokens(18_000, &assistant), 80_000); - assert_eq!(live.cumulative.input, 70_000); - assert_eq!(live.cumulative.output, 300); - assert_eq!(live.cumulative.cache_read, 60_000); - } - - #[test] - fn acp_compaction_counts_provider_context_and_cached_input() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant - .parts - .push(MessagePart::usage(60_000, 200, 20_000, 5_000, 0.0)); - let messages = vec![Message::user("question"), assistant]; - - assert_eq!(acp_context_tokens(&messages, 1_000), 85_200); - } - - #[test] - fn acp_compaction_does_not_sum_context_across_tool_requests() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant - .parts - .push(MessagePart::usage(60_000, 200, 20_000, 0, 0.0)); - assistant - .parts - .push(MessagePart::usage(10_000, 100, 40_000, 0, 0.0)); - - assert_eq!( - acp_context_tokens(&[Message::user("question"), assistant], 0), - 80_200 - ); - } - - #[test] - fn acp_compaction_uses_latest_turn_usage_and_counts_pending_history() { - use crate::session::types::{Message, MessagePart}; - let mut old = Message::assistant("old answer"); - old.parts.push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); - let mut latest = Message::assistant("new answer"); - latest - .parts - .push(MessagePart::usage(50_000, 200, 0, 0, 0.0)); - let pending = Message::user("continue"); - let pending_tokens = crate::session::compaction::message_context_tokens(&pending); - let messages = vec![Message::user("question"), old, latest, pending]; - - assert_eq!(acp_context_tokens(&messages, 0), 50_200 + pending_tokens); - } - - #[test] - fn acp_compaction_discards_provider_usage_before_compaction_boundary() { - use crate::session::types::{CompactionStats, Message, MessagePart}; - let mut assistant = Message::assistant("old answer"); - assistant - .parts - .push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); - let mut summary = Message::user(format!( - "{}\nsummary", - crate::session::compaction::SUMMARY_PREFIX - )); - summary - .parts - .push(MessagePart::usage(95_000, 500, 0, 0, 0.0)); - let marker = crate::session::compaction::compaction_marker(CompactionStats { - before_tokens: 90_100, - after_tokens: 100, - before_messages: 2, - after_messages: 1, - }); - let messages = vec![Message::user("question"), assistant, summary, marker]; - - assert_eq!( - acp_context_tokens(&messages, 1_000), - crate::session::compaction::total_context_tokens(&messages) + 1_000 - ); - } - - #[test] - fn acp_compaction_honors_reserved_tokens_and_auto_setting() { - use crate::config::configuration::CompactionConfig; - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant - .parts - .push(MessagePart::usage(70_000, 0, 10_000, 0, 0.0)); - let messages = vec![Message::user("history ".repeat(2_000)), assistant]; - assert!(crate::session::compaction::select_messages_for_compaction( - &messages, - crate::session::compaction::DEFAULT_TAIL_TURNS, - ) - .is_some()); - let config = CompactionConfig::Settings { - auto: true, - prune: false, - reserved: Some(20_000), - }; - let used = acp_context_tokens(&messages, 1_000); - assert!(crate::session::compaction::total_context_tokens(&messages) < 80_000); - assert!(crate::session::compaction::should_auto_compact( - &config, - used, - Some(100_000), - None - )); - assert!(!crate::session::compaction::should_auto_compact( - &config, - used, - Some(100_001), - None - )); - assert!(!crate::session::compaction::should_auto_compact( - &CompactionConfig::Disabled, - used, - Some(100_000), - None - )); - assert!(!crate::session::compaction::should_auto_compact( - &config, used, None, None - )); - let no_auto = CompactionConfig::Settings { - auto: false, - prune: false, - reserved: Some(20_000), - }; - assert!(!crate::session::compaction::should_auto_compact( - &no_auto, - used, - Some(100_000), - None - )); - } - - #[test] - fn acp_compaction_counts_prompt_and_pending_tokens_without_provider_usage() { - use crate::config::configuration::CompactionConfig; - use crate::session::types::Message; - let messages = vec![Message::user("history ".repeat(2_000))]; - let history_tokens = crate::session::compaction::total_context_tokens(&messages); - let config = CompactionConfig::Settings { - auto: true, - prune: false, - reserved: Some(1_000), - }; - let context_window = Some(history_tokens as u32 + 3_000); - assert!(!crate::session::compaction::should_auto_compact( - &config, - acp_context_tokens(&messages, 1_500), - context_window, - None, - )); - assert!(crate::session::compaction::should_auto_compact( - &config, - acp_context_tokens(&messages, 1_500) + 500, - context_window, - None, - )); - } - - #[test] - fn acp_compaction_replaces_previous_prompt_estimate_in_provider_context() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - let mut usage = MessagePart::usage(80_000, 200, 0, 0, 0.0); - usage.data["acp_prompt_tokens"] = serde_json::json!(5_000); - assistant.parts.push(usage); - let messages = vec![Message::user("question"), assistant]; - assert_eq!(acp_context_tokens(&messages, 25_000), 100_200); - assert_eq!(acp_context_tokens(&messages, 1_000), 76_200); - } - - #[test] - fn acp_compaction_keeps_provider_context_when_new_input_usage_is_missing() { - use crate::session::types::{Message, MessagePart}; - for output in [0, 100] { - let mut previous = Message::assistant("previous answer"); - previous - .parts - .push(MessagePart::usage(85_000, 200, 0, 0, 0.0)); - let mut latest = Message::assistant("new answer"); - let mut usage = MessagePart::usage(0, output, 0, 0, 0.0); - usage.data["acp_prompt_tokens"] = serde_json::json!(1_000); - latest.parts.push(usage); - let following_tokens = crate::session::compaction::message_context_tokens(&latest); - assert_eq!( - acp_context_tokens(&[Message::user("question"), previous, latest], 2_000), - 85_200 + following_tokens - ); - } - } - - #[test] - fn acp_compaction_counts_tool_output_after_provider_usage() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant - .parts - .push(MessagePart::usage(80_000, 200, 0, 0, 0.0)); - record_tool_result( - &mut assistant, - &crate::llm::ToolCallResult { - tool_call_id: "call_1".into(), - role: "tool".into(), - name: "read".into(), - content: serde_json::json!({"output_preview": "late tool output ".repeat(1_000)}) - .to_string(), - }, - ); - let mut following = Message::assistant(""); - following.parts = assistant.parts[2..].to_vec(); - let extra_tokens = crate::session::compaction::message_context_tokens(&following); - assert!(extra_tokens > 1_000); - assert_eq!( - acp_context_tokens(&[Message::user("question"), assistant], 0), - 80_200 + extra_tokens - ); - } - - #[test] - fn acp_compaction_keeps_fresh_system_prompt_in_request() { - use crate::session::types::{CompactionStats, Message}; - let summary = Message::user(format!( - "{}\nsummary", - crate::session::compaction::SUMMARY_PREFIX - )); - let tail = Message::assistant("retained answer"); - let marker = crate::session::compaction::compaction_marker(CompactionStats { - before_tokens: 10_000, - after_tokens: 100, - before_messages: 3, - after_messages: 2, - }); - let pending = Message::user("continue"); - let history = vec![ - Message::user("archived history"), - summary.clone(), - tail.clone(), - marker, - pending.clone(), - ]; - let system = Message::system("fresh system instructions"); - let request = acp_request_messages(&history, system.clone()); - assert_eq!(request, vec![system, summary, tail, pending]); - assert_eq!( - crate::session::compaction::filter_messages_for_context(&request), - request - ); - } - - #[test] - fn acp_compaction_cost_keeps_summary_billing_without_double_counting() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant.cost = Some(0.25); - assistant - .parts - .push(MessagePart::usage(100, 20, 0, 0, 0.25)); - let mut summary = Message::user(format!( - "{}\nsummary", - crate::session::compaction::SUMMARY_PREFIX - )); - summary.parts.push(MessagePart::usage(100, 20, 0, 0, 0.05)); - assert!((acp_session_cost(&[assistant, summary]) - 0.30).abs() < f64::EPSILON); - } - - #[tokio::test(flavor = "current_thread")] - async fn acp_compaction_notifies_client_of_reduced_context_and_summary_cost() { - use crate::session::types::{Message, MessagePart}; - use agent_client_protocol::Agent; - tokio::task::LocalSet::new() - .run_until(async { - let workspace = tempfile::tempdir().unwrap(); - let mut manager = SessionManager::new() - .with_history_for_workspace(workspace.path()) - .unwrap(); - let session_id = manager.create_session(Some("ACP compaction regression".into())); - let mut assistant = Message::assistant("old answer"); - assistant - .parts - .push(MessagePart::usage(80_000, 200, 0, 0, 0.25)); - let messages = vec![Message::user("history ".repeat(2_000)), assistant]; - let selection = crate::session::compaction::select_messages_for_compaction( - &messages, - crate::session::compaction::DEFAULT_TAIL_TURNS, - ) - .unwrap(); - let mut session = test_session(); - session.context_window = Some(100_000); - let summary = crate::llm::client::CompactionSummary { - text: "Retained task details".into(), - usage: crate::aisdk::chunk::TokenUsage { - input: 3_000, - output: 20, - cache_read: 0, - cache_write: 0, - }, - }; - let (mut compacted, stats) = compacted_messages( - &messages, - &selection, - &summary, - &session, - crate::session::compaction::total_context_tokens(&messages), - messages.len(), - ) - .unwrap(); - let summary_message = compacted - .iter_mut() - .find(|message| crate::session::compaction::is_compaction_summary(message)) - .unwrap(); - summary_message - .parts - .push(MessagePart::usage(0, 1, 0, 0, 0.05)); - let expected_used = - crate::session::compaction::total_context_tokens(&compacted) + 1_000; - manager - .replace_session_messages(&session_id, compacted) - .unwrap(); - let service = AcpService { - sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), - session_manager: Arc::new(Mutex::new(manager)), - client_capabilities: Arc::new(Mutex::new(Default::default())), - }; - let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); - let client = Client.builder().on_receive_notification( - async move |notification: SessionNotification, _connection| { - let _ = observed_tx.send(notification.update); - Ok(()) - }, - agent_client_protocol::on_receive_notification!(), - ); - Agent - .builder() - .connect_with(client, async move |connection| { - service - .notify_compaction(&session_id, &session, &connection, stats, 1_000) - .unwrap(); - let feedback = tokio::time::timeout( - std::time::Duration::from_secs(2), - observed_rx.recv(), - ) - .await - .unwrap() - .unwrap(); - assert!(matches!(feedback, SessionUpdate::AgentMessageChunk(_))); - let update = tokio::time::timeout( - std::time::Duration::from_secs(2), - observed_rx.recv(), - ) - .await - .unwrap() - .unwrap(); - let SessionUpdate::UsageUpdate(usage) = update else { - panic!("missing usage update") - }; - assert_eq!(usage.used, expected_used as u64); - assert_eq!(usage.size, 100_000); - assert!((usage.cost.unwrap().amount - 0.30).abs() < f64::EPSILON); - service - .session_manager - .lock() - .unwrap() - .try_delete_session(&session_id) - .unwrap(); - Ok(()) - }) - .await - .unwrap(); - }) - .await; - } - - #[test] - fn session_info_includes_hierarchy_metadata() { - let now = std::time::SystemTime::now(); - let info = acp_session_info( - crate::session::manager::SessionInfo { - id: "child".to_string(), - parent_id: Some("parent".to_string()), - title: "Child".to_string(), - created_at: now, - updated_at: now, - message_count: 0, - workspace_id: 1, - workspace_path: "/tmp".to_string(), - workspace_name: "tmp".to_string(), - workspace_sort_order: 0, - status: crate::session::types::SessionStatus::Idle, - pinned_at: None, - archived_at: None, - }, - Some("root".to_string()), - ); - let meta = info.meta.expect("session metadata"); - assert_eq!(meta["crabcode"]["parentSessionId"], "parent"); - assert_eq!(meta["crabcode"]["rootSessionId"], "root"); - } - - #[test] - fn session_list_cursor_pages_all_sessions() { - assert_eq!( - session_page(None, 250).unwrap(), - (0, 100, Some("100".to_string())) - ); - assert_eq!( - session_page(Some("100"), 250).unwrap(), - (100, 200, Some("200".to_string())) - ); - assert_eq!(session_page(Some("200"), 250).unwrap(), (200, 250, None)); - assert!(session_page(Some("invalid"), 250).is_err()); - } - - #[test] - fn terminal_support_tracks_client_capability() { - let service = AcpService::new(Path::new("/tmp")).unwrap(); - assert!(!service.supports_terminals()); - - service.set_client_capabilities( - agent_client_protocol::schema::v1::ClientCapabilities::new().terminal(true), - ); - assert!(service.supports_terminals()); - } - - #[test] - fn builds_titles_from_tool_input() { - assert_eq!( - tool_title("bash", &serde_json::json!({"command": "cargo test"})), - "cargo test" - ); - assert_eq!( - tool_title("read", &serde_json::json!({"filePath": "src/main.rs"})), - "src/main.rs" - ); - } - - #[test] - fn acp_tool_result_prefers_full_output() { - let payload = serde_json::json!({ - "output": "complete tool output", - "output_preview": "short preview", - }); - - assert_eq!(tool_result_text(&payload), "complete tool output"); - } - - #[test] - fn acp_tool_result_supports_legacy_preview_payloads() { - let payload = serde_json::json!({"output_preview": "legacy output"}); - - assert_eq!(tool_result_text(&payload), "legacy output"); - } - - #[test] - fn acp_tool_locations_normalize_multi_file_and_patch_paths() { - let cwd = Path::new("/tmp/workspace"); - let write_locations = tool_locations( - "write_files", - &serde_json::json!({ - "files": [ - {"file_path": "src/a.rs", "content": "a"}, - {"file_path": "/tmp/b.rs", "content": "b"} - ] - }), - cwd, - ); - assert_eq!( - write_locations[0].path, - PathBuf::from("/tmp/workspace/src/a.rs") - ); - assert_eq!(write_locations[1].path, PathBuf::from("/tmp/b.rs")); - - let patch_locations = tool_locations( - "apply_patch", - &serde_json::json!({ - "patch": "*** Begin Patch\n*** Update File: src/a.rs\n*** Add File: src/b.rs\n*** End Patch" - }), - cwd, - ); - assert_eq!(patch_locations.len(), 2); - assert_eq!( - patch_locations[1].path, - PathBuf::from("/tmp/workspace/src/b.rs") - ); - } - - #[test] - fn acp_tool_result_emits_diff_location_and_image_content() { - let payload = serde_json::json!({ - "output": "updated", - "metadata": { - "path": "src/main.rs", - "line_number": 4, - "old_text": "fn old() {}", - "new_text": "fn new() {}" - }, - "images": [{ - "data_url": "data:image/png;base64,aGk=", - "media_type": "image/png" - }] - }); - let cwd = Path::new("/tmp/workspace"); - - let locations = tool_result_locations(&payload, cwd); - assert_eq!( - locations[0].path, - PathBuf::from("/tmp/workspace/src/main.rs") - ); - assert_eq!(locations[0].line, Some(4)); - - let content = tool_result_content(&payload, cwd); - let diff = content.iter().find_map(|item| match item { - ToolCallContent::Diff(diff) => Some(diff), - _ => None, - }); - let diff = diff.expect("diff content"); - assert_eq!(diff.path, PathBuf::from("/tmp/workspace/src/main.rs")); - assert_eq!(diff.old_text.as_deref(), Some("fn old() {}")); - assert_eq!(diff.new_text, "fn new() {}"); - assert!(content.iter().any(|item| matches!( - item, - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::Image(image) if image.data == "aGk=" && image.mime_type == "image/png") - ))); - } - - #[test] - fn acp_permission_prefers_originating_tool_call_id() { - assert_eq!(permission_tool_call_id(Some("call_123")), "call_123"); - } - - #[test] - fn acp_permission_generates_fallback_id_without_origin() { - assert!(permission_tool_call_id(None).starts_with("permission:")); - } - - #[test] - fn acp_permission_offers_explicit_session_wide_allow() { - let options = permission_options(); - let always = options - .iter() - .find(|option| option.option_id.to_string() == "always") - .expect("always option"); - assert_eq!(always.name, "Always allow for this session"); - assert_eq!(always.kind, PermissionOptionKind::AllowAlways); - } - - #[tokio::test] - async fn acp_session_reuses_permission_grants_across_turns() { - let temp = tempfile::tempdir().expect("temp dir"); - let workspace = temp.path().join("workspace"); - let external = temp.path().join("external"); - std::fs::create_dir_all(&workspace).expect("workspace"); - std::fs::create_dir_all(&external).expect("external"); - - let mut session = test_session(); - session.cwd = workspace.clone(); - session.tool_permissions = crate::tools::ToolPermissions::new(&workspace); - let first_turn = tool_permissions(&session); - let second_turn = tool_permissions(&session); - let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); - let params = serde_json::json!({ "file_path": external.join("one.txt") }); - - let pending = tokio::spawn({ - let permissions = first_turn.clone(); - let params = params.clone(); - let tx = tx.clone(); - async move { - permissions - .preflight("build", "read", ¶ms, Some(&tx)) - .await - } - }); - let prompt = match rx.recv().await { - Some(crate::llm::ChunkMessage::PermissionRequest(prompt)) => prompt, - _ => panic!("expected permission request"), - }; - let _ = prompt - .response_tx - .send(crate::tools::PermissionResponse::AllowAlways); - assert!(pending.await.expect("permission task").is_ok()); - - let next_params = serde_json::json!({ "file_path": external.join("nested/two.txt") }); - assert!(second_turn - .preflight("build", "read", &next_params, Some(&tx)) - .await - .is_ok()); - assert!(rx.try_recv().is_err()); - } - - #[test] - fn acp_permission_edit_includes_preflight_diff() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("main.rs"); - std::fs::write(&path, "fn old() {}\n").unwrap(); - let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); - let prompt = crate::tools::PermissionPrompt { - tool_call_id: Some("call_edit".to_string()), - tool_id: "edit".to_string(), - action: crate::tools::PermissionAction::Write, - permission: "edit".to_string(), - patterns: vec![path.to_string_lossy().into_owned()], - target: Some(path.to_string_lossy().into_owned()), - command: None, - workdir: None, - workspace: dir.path().to_string_lossy().into_owned(), - reason: "approval".to_string(), - raw_input: serde_json::json!({ - "file_path": path, - "old_string": "old", - "new_string": "new", - "replace_all": false - }), - response_tx, - }; - - let content = permission_tool_content(&prompt, dir.path()); - let ToolCallContent::Diff(diff) = &content[0] else { - panic!("expected preflight diff"); - }; - assert_eq!(diff.old_text.as_deref(), Some("fn old() {}\n")); - assert_eq!(diff.new_text, "fn new() {}\n"); - } - - #[test] - fn acp_permission_apply_patch_includes_preflight_diff() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("note.txt"); - std::fs::write(&path, "before\n").unwrap(); - let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); - let prompt = crate::tools::PermissionPrompt { - tool_call_id: Some("call_patch".to_string()), - tool_id: "apply_patch".to_string(), - action: crate::tools::PermissionAction::Write, - permission: "edit".to_string(), - patterns: vec![path.to_string_lossy().into_owned()], - target: Some(path.to_string_lossy().into_owned()), - command: None, - workdir: None, - workspace: dir.path().to_string_lossy().into_owned(), - reason: "approval".to_string(), - raw_input: serde_json::json!({ - "patch": format!( - "*** Begin Patch\n*** Update File: {}\n@@\n-before\n+after\n*** End Patch\n", - path.display() - ) - }), - response_tx, - }; - - let content = permission_tool_content(&prompt, dir.path()); - let ToolCallContent::Diff(diff) = &content[0] else { - panic!("expected apply_patch preflight diff"); - }; - assert_eq!(diff.path, path); - assert_eq!(diff.old_text.as_deref(), Some("before\n")); - assert_eq!(diff.new_text, "after\n"); - } - - #[test] - fn acp_question_form_preserves_single_multi_custom_and_scope() { - let form = acp_question_form( - "session_1", - Some("question_call_1"), - &serde_json::json!([ - { - "question": "Pick one", - "header": "Single", - "options": [ - {"label": "A", "description": "First"}, - {"label": "B", "description": "Second"} - ] - }, - { - "question": "Pick several", - "header": "Multiple", - "multiple": true, - "options": [ - {"label": "X", "description": "First"}, - {"label": "Y", "description": "Second"} - ] - } - ]), - ); - let wire = serde_json::to_value(&form.request).expect("elicitation request"); - - assert_eq!(wire["mode"], "form"); - assert_eq!(wire["sessionId"], "session_1"); - assert_eq!(wire["toolCallId"], "question_call_1"); - assert_eq!( - wire["requestedSchema"]["properties"]["question_0"]["type"], - "string" - ); - assert_eq!( - wire["requestedSchema"]["properties"]["question_0"]["oneOf"][0]["title"], - "A" - ); - assert_eq!( - wire["requestedSchema"]["properties"]["question_1"]["type"], - "array" - ); - assert_eq!( - wire["requestedSchema"]["properties"]["question_1_custom"]["type"], - "string" - ); - } - - #[test] - fn acp_question_answers_restore_labels_and_custom_text() { - let form = acp_question_form( - "session_1", - None, - &serde_json::json!([ - { - "question": "Pick one", - "options": [{"label": "A"}, {"label": "B"}] - }, - { - "question": "Pick several", - "multiple": true, - "options": [{"label": "X"}, {"label": "Y"}] - } - ]), - ); - let mut content = std::collections::BTreeMap::new(); - content.insert( - "question_0".to_string(), - ElicitationContentValue::String("q0_option_1".to_string()), - ); - content.insert( - "question_1".to_string(), - ElicitationContentValue::StringArray(vec![ - "q1_option_0".to_string(), - "q1_option_1".to_string(), - ]), - ); - content.insert( - "question_1_custom".to_string(), - ElicitationContentValue::String("Other choice".to_string()), - ); - let action = ElicitationAction::Accept( - agent_client_protocol::schema::v1::ElicitationAcceptAction::new().content(content), - ); - - assert_eq!( - acp_question_answers(&form.fields, action), - serde_json::json!([["B"], ["X", "Y", "Other choice"]]) - ); - assert_eq!( - acp_question_answers(&form.fields, ElicitationAction::Cancel), - serde_json::json!([[], []]) - ); - - let mut custom_content = std::collections::BTreeMap::new(); - custom_content.insert( - "question_0".to_string(), - ElicitationContentValue::String("q0_option_0".to_string()), - ); - custom_content.insert( - "question_0_custom".to_string(), - ElicitationContentValue::String("Custom only".to_string()), - ); - let custom_action = ElicitationAction::Accept( - agent_client_protocol::schema::v1::ElicitationAcceptAction::new() - .content(custom_content), - ); - assert_eq!( - acp_question_answers(&form.fields, custom_action), - serde_json::json!([["Custom only"], []]) - ); - } - - #[test] - fn maps_typed_turn_stop_reasons_to_acp() { - assert_eq!( - acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTokens)), - StopReason::MaxTokens - ); - assert_eq!( - acp_stop_reason(Some(crate::llm::TurnStopReason::Refusal)), - StopReason::Refusal - ); - assert_eq!( - acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), - StopReason::MaxTurnRequests - ); - assert_eq!( - acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), - StopReason::MaxTurnRequests - ); - assert_eq!(acp_stop_reason(None), StopReason::EndTurn); - } - - #[test] - fn recognizes_only_exact_compact_control_command() { - assert_eq!(compact_command("/compact").unwrap(), true); - assert_eq!(compact_command(" /compact ").unwrap(), true); - assert_eq!(compact_command("/compactness").unwrap(), false); - assert_eq!(compact_command("hello").unwrap(), false); - assert!(compact_command("/compact extra").is_err()); - } - - #[test] - fn advertises_compact_as_no_input_command() { - let session = test_session(); - let command = available_commands(&session) - .into_iter() - .find(|command| command.name == "compact") - .expect("compact command"); - assert_eq!( - command.description, - "Summarize this session to reduce context" - ); - assert!(command.input.is_none()); - } - - #[tokio::test] - async fn advertises_and_parses_btw_side_command() { - let session = test_session(); - let command = available_commands(&session) - .into_iter() - .find(|command| command.name == "btw") - .expect("btw command"); - assert!(command.input.is_some()); - assert_eq!( - expand_slash_command(&session, "/btw what changed?") - .await - .expect("btw expansion"), - SlashExpansion::Btw("what changed?".to_string()) - ); - assert!(expand_slash_command(&session, "/btw").await.is_err()); - } - - #[test] - fn builds_smaller_soft_compaction_for_acp() { - let session = test_session(); - let messages = vec![ - crate::session::types::Message::user("u".repeat(8_000)), - crate::session::types::Message::assistant("a".repeat(8_000)), - crate::session::types::Message::user("recent"), - ]; - let selection = crate::session::compaction::select_messages_for_compaction_with_min( - &messages, - crate::session::compaction::DEFAULT_TAIL_TURNS, - 0, - ) - .expect("compaction selection"); - let before_tokens = crate::session::compaction::total_context_tokens(&messages); - let before_messages = - crate::session::compaction::filter_messages_for_context(&messages).len(); - - let (compacted, stats) = compacted_messages( - &messages, - &selection, - &crate::llm::client::CompactionSummary { - text: "short handoff".to_string(), - usage: Default::default(), - }, - &session, - before_tokens, - before_messages, - ) - .expect("smaller compaction"); - - assert!(stats.after_tokens < stats.before_tokens); - assert!(crate::session::compaction::latest_compaction_stats(&compacted).is_some()); - assert_eq!(compacted[0].id, messages[0].id); - } - - #[test] - fn flattens_text_and_embedded_context() { - let text = prompt_text(vec![ - ContentBlock::from("Inspect this."), - ContentBlock::Resource(agent_client_protocol::schema::v1::EmbeddedResource::new( - EmbeddedResourceResource::TextResourceContents( - agent_client_protocol::schema::v1::TextResourceContents::new( - "fn main() {}", - "file:///tmp/main.rs", - ), - ), - )), - ]) - .expect("prompt text"); - - assert_eq!(text, "Inspect this.[file:///tmp/main.rs]\nfn main() {}"); - } - - #[test] - fn writes_supported_acp_image_to_managed_session_storage() { - let session_id = format!("acp-image-{}", cuid2::create_id()); - let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/png"); - let path = write_prompt_image(&session_id, &image).expect("image file"); - - assert!(Path::new(&path).starts_with(crate::persistence::attachments::root_dir())); - assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); - crate::persistence::attachments::cleanup_session(&session_id).unwrap(); - } - - #[test] - fn command_text_override_preserves_embedded_context() { - let resource = agent_client_protocol::schema::v1::EmbeddedResource::new( - EmbeddedResourceResource::TextResourceContents( - agent_client_protocol::schema::v1::TextResourceContents::new( - "important context", - "file:///tmp/context.txt", - ), - ), - ); - let (text, images, audio) = prompt_content( - vec![ - ContentBlock::Text(agent_client_protocol::schema::v1::TextContent::new( - "/review src/lib.rs", - )), - ContentBlock::Text(agent_client_protocol::schema::v1::TextContent::new( - "\nAdditional instructions", - )), - ContentBlock::Resource(resource), - ], - Some("expanded command prompt"), - false, - false, - "test-session", - &test_session(), - ) - .expect("prompt content"); - - assert_eq!( - text, - "expanded command prompt\nAdditional instructions[file:///tmp/context.txt]\nimportant context" - ); - assert!(images.is_empty()); - assert!(audio.is_empty()); - } - - #[test] - fn writes_supported_acp_audio_to_managed_session_storage() { - let session_id = format!("acp-audio-{}", cuid2::create_id()); - let audio = agent_client_protocol::schema::v1::AudioContent::new("YXVkaW8=", "audio/wav"); - let path = write_prompt_audio(&session_id, &audio).expect("audio file"); - - assert!(path.ends_with(".wav")); - assert_eq!(std::fs::read(&path).unwrap(), b"audio"); - crate::persistence::attachments::cleanup_session(&session_id).unwrap(); - } - - #[test] - fn rejects_audio_for_models_without_audio_modality() { - let session_id = format!("acp-audio-{}", cuid2::create_id()); - let result = prompt_content( - vec![ContentBlock::Audio( - agent_client_protocol::schema::v1::AudioContent::new("YXVkaW8=", "audio/wav"), - )], - None, - false, - false, - &session_id, - &test_session(), - ); - - assert!(result.is_err()); - crate::persistence::attachments::cleanup_session(&session_id).unwrap(); - } - - #[test] - fn writes_acp_clipboard_image_data_uri_to_managed_storage() { - let session_id = format!("acp-image-{}", cuid2::create_id()); - let image = agent_client_protocol::schema::v1::ImageContent::new( - "data:image/png;base64,aGk=", - "application/octet-stream", - ); - let path = write_prompt_image(&session_id, &image).expect("image file"); - - assert!(path.ends_with(".png")); - assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); - crate::persistence::attachments::cleanup_session(&session_id).unwrap(); - } - - #[test] - fn prompt_image_failure_rolls_back_prior_managed_files() { - let session_id = format!("acp-image-{}", cuid2::create_id()); - let result = prompt_content( - vec![ - ContentBlock::Image(agent_client_protocol::schema::v1::ImageContent::new( - "aGk=", - "image/png", - )), - ContentBlock::Image(agent_client_protocol::schema::v1::ImageContent::new( - "not-base64", - "image/png", - )), - ], - None, - true, - false, - &session_id, - &test_session(), - ); - - assert!(result.is_err()); - let directory = crate::persistence::attachments::session_dir(&session_id).unwrap(); - assert!( - !directory.exists() - || std::fs::read_dir(&directory) - .unwrap() - .all(|entry| entry.is_err()) - ); - crate::persistence::attachments::cleanup_session(&session_id).unwrap(); - } - - #[test] - fn rejects_non_base64_acp_image_data_uri() { - let image = agent_client_protocol::schema::v1::ImageContent::new( - "data:image/png,not-base64", - "image/png", - ); - - assert!(write_prompt_image("test", &image).is_err()); - } - - #[test] - fn rejects_unsupported_acp_image_mime_type() { - let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/tiff"); - - assert!(write_prompt_image("test", &image).is_err()); - } - - fn config_with_command(command: crate::command::custom::CustomCommand) -> LoadedConfig { - let mut merged_config = crate::config::configuration::MergedConfig::default(); - merged_config.commands.push(command); - config_with_merged(merged_config) - } - - fn config_with_merged( - merged_config: crate::config::configuration::MergedConfig, - ) -> LoadedConfig { - LoadedConfig { - merged_config, - raw_merged: serde_json::Value::Null, - diagnostics: Default::default(), - inventory: Default::default(), - project_root: PathBuf::from("/tmp"), - cwd: PathBuf::from("/tmp"), - xdg_config_home: PathBuf::from("/tmp"), - } - } - - fn empty_config() -> LoadedConfig { - config_with_merged(crate::config::configuration::MergedConfig::default()) - } - - fn session_with_config(config: LoadedConfig) -> AcpSession { - let skills = crate::skill::SkillStore::load(&config.xdg_config_home, &config.project_root); - AcpSession { - cwd: config.cwd.clone(), - tool_permissions: configured_tool_permissions(&config.cwd, &config), - config, - skills, - models: Vec::new(), - provider: String::new(), - model: String::new(), - agent: "Build".to_string(), - reasoning_selection: crate::model::reasoning::ReasoningEffort::None, - reasoning: None, - context_window: None, - cancellation: None, - } - } - - #[tokio::test(flavor = "current_thread")] - async fn cancelled_stream_persists_partial_transcript_and_interrupted_status() { - cancelled_stream_finalizes(false, false).await; - } - - #[tokio::test(flavor = "current_thread")] - async fn cancelled_permission_persists_partial_transcript_and_interrupted_status() { - cancelled_stream_finalizes(true, false).await; - } - - #[tokio::test(flavor = "current_thread")] - async fn close_reopen_waits_for_old_prompt_finalization() { - cancelled_stream_finalizes(false, true).await; - } - - async fn cancelled_stream_finalizes(pending_permission: bool, close_reopen: bool) { - use agent_client_protocol::schema::v1::RequestPermissionResponse; - use agent_client_protocol::{Agent, Responder}; - use std::time::Duration; - - tokio::task::LocalSet::new().run_until(async { - let workspace = tempfile::tempdir().unwrap(); - let mut manager = SessionManager::new() - .with_history_for_workspace(workspace.path()).unwrap(); - let session_id = manager.create_session(Some("ACP cancellation regression".into())); - let db_id = manager.get_db_id(&session_id).unwrap(); - manager.set_session_status( - &session_id, crate::session::types::SessionStatus::Streaming, None, - ).unwrap(); - let service = AcpService { - sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), - session_manager: Arc::new(Mutex::new(manager)), - client_capabilities: Arc::new(Mutex::new(Default::default())), - }; - let session = session_with_config(empty_config()); - service.sessions.lock().await.attachments.insert(session_id.clone(), session.clone()); - let cancellation = service.sessions.lock().await.begin_prompt(&session_id).unwrap(); - let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); - let (permission_tx, mut permission_rx) = tokio::sync::mpsc::unbounded_channel(); - let client = Client.builder().on_receive_request( - async move |_request: RequestPermissionRequest, - responder: Responder, _connection| { - permission_tx.send(responder).unwrap(); - Ok(()) - }, - agent_client_protocol::on_receive_request!(), - ).on_receive_notification( - async move |notification: SessionNotification, _connection| { - if matches!(notification.update, SessionUpdate::AgentMessageChunk(_)) { - let _ = observed_tx.send(()); - } - Ok(()) - }, - agent_client_protocol::on_receive_notification!(), - ); - Agent.builder().connect_with(client, async move |connection| { - let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); - sender.send(crate::llm::ChunkMessage::Text("partial answer".into())).unwrap(); - sender.send(crate::llm::ChunkMessage::Reasoning("partial reasoning".into())).unwrap(); - let (response_tx, response_rx) = tokio::sync::oneshot::channel(); - if pending_permission { - sender.send(crate::llm::ChunkMessage::PermissionRequest(crate::tools::PermissionPrompt { - tool_call_id: Some("pending-tool".into()), - tool_id: "bash".into(), - action: crate::tools::PermissionAction::Write, - permission: "bash".into(), - patterns: Vec::new(), - target: None, - command: Some("mock-command".into()), - workdir: None, - workspace: workspace.path().to_string_lossy().into_owned(), - reason: "regression test".into(), - raw_input: serde_json::json!({"command": "mock-command"}), - response_tx, - })).unwrap(); - } - let stream = service.receive_prompt_stream( - &session_id, &session, &connection, &cancellation, receiver, 0, 0, 0.0, - ); - tokio::pin!(stream); - tokio::time::timeout(Duration::from_secs(2), async { - tokio::select! { - _ = observed_rx.recv() => {}, - result = &mut stream => panic!("stream ended before cancellation: {result:?}"), - } - }).await.unwrap(); - let _pending_responder = if pending_permission { - Some(tokio::time::timeout(Duration::from_secs(2), async { - tokio::select! { - responder = permission_rx.recv() => responder.unwrap(), - result = &mut stream => panic!("stream ended before permission: {result:?}"), - } - }).await.unwrap()) - } else { - None - }; - if close_reopen { - service.close_session(&session_id).await; - assert!(cancellation.is_cancelled()); - assert!(!service.sessions.lock().await.attachments.contains_key(&session_id)); - assert!(service.available_commands(&session_id).await.is_err()); - assert!(service.set_mode(&session_id, "Build").await.is_err()); - assert!(service.set_model(&session_id, "unused").await.is_err()); - assert!(service.set_reasoning_effort(&session_id, "none").await.is_err()); - // The same guard serves load and resume. Resume must fail before - // any provider discovery, while the old stream still owns writes. - assert!(service.resume_session(session_id.clone(), workspace.path().into()).await.is_err()); - assert!(service.load_session(session_id.clone(), workspace.path().into(), connection.clone()).await.is_err()); - assert!(service.sessions.lock().await.begin_prompt(&session_id).is_err()); - } else { - cancellation.cancel(); - } - let response = tokio::time::timeout(Duration::from_secs(1), &mut stream) - .await.unwrap().unwrap(); - assert_eq!(response.stop_reason, StopReason::Cancelled); - if pending_permission { - assert!(matches!(response_rx.await.unwrap(), crate::tools::PermissionResponse::Deny)); - } - let manager = service.session_manager.lock().unwrap(); - let stored = manager.get_session_ref(&session_id).unwrap(); - assert_eq!(stored.status, crate::session::types::SessionStatus::Interrupted); - assert_eq!(stored.messages.len(), 1); - let assistant = &stored.messages[0]; - assert_eq!(assistant.content, "partial answer"); - assert!(assistant.is_complete); - assert!(assistant.was_interrupted); - drop(manager); - // A fresh DAO verifies the real database, not just the in-memory message. - let dao = crate::persistence::HistoryDAO::new_for_workspace(workspace.path()).unwrap(); - assert_eq!(dao.get_session(db_id).unwrap().unwrap().status, "interrupted"); - let messages = dao.get_messages(db_id).unwrap(); - assert_eq!(messages.len(), 1); - assert!(messages[0].parts.iter().any(|part| - part.part_type == "text" && part.data["text"] == "partial answer" - )); - drop(dao); - if close_reopen { - // Stream persistence has ended, but wrapper finalization has not: - // reopen must stay blocked through that second status write too. - assert!(service.resume_session(session_id.clone(), workspace.path().into()).await.is_err()); - service.session_manager.lock().unwrap().set_session_status( - &session_id, crate::session::types::SessionStatus::Interrupted, None, - ).unwrap(); - service.sessions.lock().await.finish_prompt(&session_id); - let mut sessions = service.sessions.lock().await; - sessions.ensure_attachable(&session_id).unwrap(); - sessions.attachments.insert(session_id.clone(), session.clone()); - let next = sessions.begin_prompt(&session_id).unwrap(); - drop(sessions); - service.session_manager.lock().unwrap().set_session_status( - &session_id, crate::session::types::SessionStatus::Streaming, None, - ).unwrap(); - assert!(!next.is_cancelled()); - assert!(service.sessions.lock().await.active_prompts.contains(&session_id)); - let dao = crate::persistence::HistoryDAO::new_for_workspace(workspace.path()).unwrap(); - assert_eq!(dao.get_session(db_id).unwrap().unwrap().status, "streaming"); - service.cancel_session(&session_id).await; - assert!(next.is_cancelled(), "the new turn retains its cancellation token"); - service.sessions.lock().await.finish_prompt(&session_id); - } else { - service.sessions.lock().await.finish_prompt(&session_id); - } - service.session_manager.lock().unwrap().try_delete_session(&session_id).unwrap(); - Ok(()) - }).await.unwrap(); - }).await; - } - - #[test] - fn advertises_custom_commands_with_argument_input() { - let config = config_with_command(crate::command::custom::CustomCommand { - name: "review".to_string(), - description: Some("Review selected code".to_string()), - template: "Review: $ARGUMENTS".to_string(), - agent: None, - model: None, - subtask: Some(false), - source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( - "/tmp/opencode.jsonc", - )), - workdir: PathBuf::from("/tmp"), - }); - - let session = session_with_config(config); - let commands = available_commands(&session); - - let command = commands - .iter() - .find(|command| command.name == "review") - .expect("review command"); - assert!(commands.iter().any(|command| command.name == "skills")); - assert!(commands.iter().any(|command| command.name == "mcp")); - assert_eq!(command.description, "Review selected code"); - assert!(matches!( - command.input, - Some(AvailableCommandInput::Unstructured(_)) - )); - } - - #[test] - fn acp_tool_result_restores_native_mcp_resources() { - let metadata = serde_json::json!({ - "mcp_result": { - "content": [ - { - "type": "image", - "data": "aGk=", - "mimeType": "image/png", - "annotations": { "priority": 0.5 } - }, - { - "type": "resource_link", - "uri": "file:///tmp/readme.md", - "name": "readme", - "mimeType": "text/markdown", - "annotations": { "priority": 0.8 } - } - ] - } - }); - - let content = mcp_native_tool_content(&metadata); - assert_eq!(content.len(), 2); - assert!(matches!( - &content[0], - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::Image(image) - if image.data == "aGk=" && image.annotations.is_some()) - )); - assert!(matches!( - &content[1], - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::ResourceLink(link) - if link.uri == "file:///tmp/readme.md") - )); - } - - #[test] - fn acp_tool_result_restores_mcp_audio_and_preserves_unknown_blocks() { - let metadata = serde_json::json!({ - "mcp_result": { - "content": [ - { "type": "audio", "data": "YXVkaW8=", "mimeType": "audio/wav" }, - { "type": "future_media", "payload": { "value": 1 } } - ] - } - }); - let content = mcp_native_tool_content(&metadata); - assert!(matches!( - &content[0], - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::Audio(audio) - if audio.data == "YXVkaW8=" && audio.mime_type == "audio/wav") - )); - assert!(matches!( - &content[1], - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::Text(text) - if text.text.contains("future_media") && text.text.contains("payload")) - )); - } - - #[tokio::test] - async fn expands_custom_slash_command_before_prompting() { - let config = config_with_command(crate::command::custom::CustomCommand { - name: "review".to_string(), - description: None, - template: "Review this carefully: $ARGUMENTS".to_string(), - agent: Some("build".to_string()), - model: Some("openai/gpt-5".to_string()), - subtask: Some(false), - source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( - "/tmp/opencode.jsonc", - )), - workdir: PathBuf::from("/tmp"), - }); - - let session = session_with_config(config); - let prompt = expand_slash_command(&session, "/review src/acp/service.rs") - .await - .expect("expanded command"); - - assert_eq!( - prompt, - SlashExpansion::Prompt { - prompt: "Review this carefully: src/acp/service.rs".to_string(), - agent: Some("build".to_string()), - model: Some("openai/gpt-5".to_string()), - } - ); - } - - #[tokio::test] - async fn advertises_and_expands_workspace_skills() { - let temp = tempfile::tempdir().expect("temp dir"); - let skill_dir = temp.path().join(".opencode/skill/reviewer"); - std::fs::create_dir_all(&skill_dir).expect("skill dir"); - std::fs::write( - skill_dir.join("SKILL.md"), - "---\nname: reviewer\ndescription: Review code carefully\n---\nInspect correctness and risks.", - ) - .expect("skill file"); - let mut config = config_with_command(crate::command::custom::CustomCommand { - name: "custom".to_string(), - description: None, - template: "$ARGUMENTS".to_string(), - agent: None, - model: None, - subtask: Some(false), - source: crate::command::custom::CustomCommandSource::Config( - temp.path().join("opencode.jsonc"), - ), - workdir: temp.path().to_path_buf(), - }); - config.project_root = temp.path().to_path_buf(); - config.cwd = temp.path().to_path_buf(); - config.xdg_config_home = temp.path().join("config"); - let session = session_with_config(config); - - let commands = available_commands(&session); - assert!(commands.iter().any(|command| command.name == "reviewer")); - let prompt = expand_slash_command(&session, "/reviewer src/lib.rs") - .await - .expect("expanded skill"); - let SlashExpansion::Prompt { prompt, .. } = prompt else { - panic!("expected skill prompt"); - }; - assert!(prompt.contains("Inspect correctness and risks.")); - assert!(prompt.contains("src/lib.rs")); - } - - #[tokio::test] - async fn reports_configured_mcp_servers() { - let mut config = config_with_command(crate::command::custom::CustomCommand { - name: "custom".to_string(), - description: None, - template: "$ARGUMENTS".to_string(), - agent: None, - model: None, - subtask: Some(false), - source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( - "/tmp/opencode.jsonc", - )), - workdir: PathBuf::from("/tmp"), - }); - config.merged_config.mcp.insert( - "filesystem".to_string(), - crate::config::configuration::McpServerConfig::Local( - crate::config::configuration::McpLocalConfig { - command: vec!["mcp-server".to_string()], - cwd: None, - environment: Default::default(), - enabled: true, - timeout_ms: None, - }, - ), - ); - let session = session_with_config(config); - - let prompt = expand_slash_command(&session, "/mcp") - .await - .expect("mcp status"); - let SlashExpansion::LocalResult(prompt) = prompt else { - panic!("expected local MCP result"); - }; - assert!(prompt.contains("filesystem (local, connecting)")); - } - - #[test] - fn merges_client_mcp_servers_into_session_config() { - let mut config = config_with_command(crate::command::custom::CustomCommand { - name: "custom".to_string(), - description: None, - template: "$ARGUMENTS".to_string(), - agent: None, - model: None, - subtask: Some(false), - source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( - "/tmp/opencode.jsonc", - )), - workdir: PathBuf::from("/tmp"), - }); - merge_acp_mcp_servers( - &mut config, - vec![serde_json::from_value(serde_json::json!({ - "name": "zed-fs", - "command": "mcp-server", - "args": ["--stdio"], - "env": [{"name": "ROOT", "value": "/workspace"}] - })) - .expect("ACP stdio MCP server")], - ); - - let server = config - .merged_config - .mcp - .get("zed-fs") - .expect("merged ACP MCP server"); - assert_eq!(server.kind(), "local"); - assert!(server.enabled()); - } - - #[test] - fn rejects_non_absolute_workspaces() { - assert!(workspace_path(Path::new("relative")).is_err()); - } - - #[test] - fn compaction_mode_descriptions_preserve_configuration_and_visibility() { - let mut session = test_session(); - assert!(!session_modes(&session) - .available_modes - .iter() - .any(|mode| mode.id.to_string() == "compaction")); - - for description in ["", " \n ", "Use this agent to preserve project decisions."] { - let definitions = crate::agent::definition::parse_agent_definitions_from_config( - Some(&serde_json::json!({"compaction": { - "mode": "primary", "description": description - }})), - &mut Vec::new(), - ); - session.config.merged_config.agent_registry = - crate::agent::definition::AgentRegistry::with_definitions(None, definitions); - let expected = if description.trim().is_empty() { - "Compaction agent. Summarizes conversation context for continued work." - } else { - description - }; - let modes = session_modes(&session); - let compaction = modes - .available_modes - .iter() - .find(|mode| mode.id.to_string() == "compaction") - .expect("configured compaction mode"); - assert_eq!(compaction.description.as_deref(), Some(expected)); - let config = serde_json::to_value(session_config_options(&session)).unwrap(); - let compaction = config[0]["options"] - .as_array() - .unwrap() - .iter() - .find(|option| option["value"] == "compaction") - .expect("configured compaction option"); - assert_eq!(compaction["description"], expected); - assert_eq!( - modes - .available_modes - .iter() - .find(|mode| mode.id.to_string() == "build") - .unwrap() - .description - .as_deref(), - Some("The default agent. Executes tools based on configured permissions.") - ); - } - - for settings in [ - serde_json::json!({"mode": "subagent"}), - serde_json::json!({"mode": "primary", "hidden": true}), - ] { - let definitions = crate::agent::definition::parse_agent_definitions_from_config( - Some(&serde_json::json!({"compaction": settings})), - &mut Vec::new(), - ); - session.config.merged_config.agent_registry = - crate::agent::definition::AgentRegistry::with_definitions(None, definitions); - assert!(!session_modes(&session) - .available_modes - .iter() - .any(|mode| mode.id.to_string() == "compaction")); - let config = serde_json::to_value(session_config_options(&session)).unwrap(); - assert!(!config[0]["options"] - .as_array() - .unwrap() - .iter() - .any(|option| option["value"] == "compaction")); - } - } - - #[test] - fn builds_grouped_model_config_option() { - let option = model_config_option( - &[ - model("anthropic", "Anthropic", "claude", "Claude"), - model("openai", "OpenAI", "gpt-5", "GPT-5"), - model("openai", "OpenAI", "gpt-5-mini", "GPT-5 Mini"), - ], - "openai", - "gpt-5", - ); - - assert_eq!(option.id.to_string(), "model"); - assert_eq!(option.category, Some(SessionConfigOptionCategory::Model)); - let agent_client_protocol::schema::v1::SessionConfigKind::Select(select) = option.kind - else { - panic!("model option should be a select"); - }; - assert_eq!(select.current_value.to_string(), "openai/gpt-5"); - let agent_client_protocol::schema::v1::SessionConfigSelectOptions::Grouped(groups) = - select.options - else { - panic!("model options should be grouped"); - }; - assert_eq!(groups.len(), 2); - assert_eq!(groups[1].name, "OpenAI"); - assert_eq!(groups[1].options[1].value.to_string(), "openai/gpt-5-mini"); - } - - #[test] - fn validates_model_refs_against_selectable_catalog() { - let models = [model("openai", "OpenAI", "gpt-5", "GPT-5")]; - - assert_eq!( - find_selectable_model(&models, "openai/gpt-5") - .expect("known model") - .id, - "gpt-5" - ); - assert!(find_selectable_model(&models, "openai/missing").is_err()); - assert!(find_selectable_model(&models, "other/gpt-5").is_err()); - } - - #[test] - fn resolves_context_window_from_selectable_models() { - let mut model = model("example", "Example", "large-context", "Large Context"); - model.context_window = Some(1_090_000); - - assert_eq!( - model_context_window(&empty_config(), &[model], "example", "large-context"), - Some(1_090_000) - ); - } - - #[test] - fn preserves_selected_reasoning_effort_when_model_cannot_apply_it() { - let session = test_session(); - let option = reasoning_config_option(&session); - assert_eq!(option.id.to_string(), "effort"); - assert_eq!(option.name, "Effort"); - assert_eq!( - option.category, - Some(SessionConfigOptionCategory::ThoughtLevel) - ); - let agent_client_protocol::schema::v1::SessionConfigKind::Select(select) = option.kind - else { - panic!("reasoning option should be a select"); - }; - assert_eq!(select.current_value.to_string(), "high"); - let agent_client_protocol::schema::v1::SessionConfigSelectOptions::Ungrouped(options) = - select.options - else { - panic!("reasoning options should be ungrouped"); - }; - let values = options - .iter() - .map(|option| (option.value.to_string(), option.name.clone())) - .collect::>(); - assert_eq!( - values, - vec![ - ("none".to_string(), "None".to_string()), - ("low".to_string(), "Low".to_string()), - ("medium".to_string(), "Medium".to_string()), - ("high".to_string(), "High".to_string()), - ("xhigh".to_string(), "Xhigh".to_string()), - ("max".to_string(), "Max".to_string()), - ] - ); - } - - #[test] - fn unknown_model_usage_cost_is_zero() { - let mut session = session_with_config(LoadedConfig { - merged_config: crate::config::configuration::MergedConfig::default(), - raw_merged: serde_json::Value::Null, - diagnostics: Default::default(), - inventory: Default::default(), - project_root: PathBuf::from("/tmp"), - cwd: PathBuf::from("/tmp"), - xdg_config_home: PathBuf::from("/tmp"), - }); - session.provider = "no-such-provider".to_string(); - session.model = "no-such-model".to_string(); - let usage = crate::aisdk::chunk::TokenUsage { - input: 1_000, - output: 1_000, - cache_read: 0, - cache_write: 0, - }; - assert_eq!(estimate_session_usage_cost(&session, &usage), 0.0); - } -} - -#[cfg(test)] -mod cancellation_lifecycle_tests { - use super::*; - - async fn delayed_terminal_creation_cleanup(acknowledge_kill: bool) { - use agent_client_protocol::schema::v1::{ - CreateTerminalResponse, KillTerminalResponse, ReleaseTerminalResponse, - }; - use agent_client_protocol::{Agent, Responder}; - use std::time::Duration; - - tokio::task::LocalSet::new() - .run_until(async { - let cancellation = CancellationToken::new(); - let (create_tx, mut create_rx) = tokio::sync::mpsc::unbounded_channel(); - let (cleanup_tx, mut cleanup_rx) = tokio::sync::mpsc::unbounded_channel(); - let kill_tx = cleanup_tx.clone(); - let (pending_kill_tx, mut pending_kill_rx) = - tokio::sync::mpsc::unbounded_channel(); - let client = Client - .builder() - .on_receive_request( - async move |request: CreateTerminalRequest, - responder: Responder, - _cx: ConnectionTo| { - assert_eq!(request.session_id.to_string(), "cancelled-session"); - assert_eq!(request.command, "bash"); - // Retain the response until the prompt-side bridge has returned. - create_tx.send(responder).unwrap(); - Ok(()) - }, - agent_client_protocol::on_receive_request!(), - ) - .on_receive_request( - async move |request: KillTerminalRequest, - responder: Responder, - _cx: ConnectionTo| { - kill_tx - .send(("kill", request.session_id, request.terminal_id)) - .unwrap(); - if acknowledge_kill { - responder.respond(KillTerminalResponse::new()) - } else { - // Handlers run in the client's dispatch loop: waiting here - // would prevent it from dispatching terminal/release at all. - // Retain the responder without replying, but let dispatch - // continue so release must arrive with kill still outstanding. - pending_kill_tx.send(responder).unwrap(); - Ok(()) - } - }, - agent_client_protocol::on_receive_request!(), - ) - .on_receive_request( - async move |request: ReleaseTerminalRequest, - responder: Responder, - _cx: ConnectionTo| { - cleanup_tx - .send(("release", request.session_id, request.terminal_id)) - .unwrap(); - responder.respond(ReleaseTerminalResponse::new()) - }, - agent_client_protocol::on_receive_request!(), - ); - - let result = Agent - .builder() - .connect_with(client, async move |connection| { - let (control_tx, mut control_rx) = - tokio::sync::mpsc::unbounded_channel(); - let request = crate::tools::TerminalSessionRequest { - start: crate::tools::TerminalSessionStart { - session_id: "cancelled-session".to_string(), - tool_call_id: "terminal-call".to_string(), - command: "mock-command-never-executed".to_string(), - description: "Delayed mock terminal".to_string(), - workdir: None, - cols: 80, - rows: 24, - job_id: None, - }, - control_tx, - }; - let bridge = bridge_terminal_session( - &connection, - "cancelled-session", - Path::new("/tmp"), - request, - &cancellation, - ); - tokio::pin!(bridge); - let responder = tokio::time::timeout(Duration::from_secs(2), async { - tokio::select! { - responder = create_rx.recv() => responder.unwrap(), - _ = &mut bridge => panic!("bridge returned before creation or cancellation"), - } - }) - .await - .expect("mock client did not receive terminal/create"); - - cancellation.cancel(); - tokio::time::timeout(Duration::from_secs(1), &mut bridge) - .await - .expect("prompt-side bridge waited for delayed terminal creation"); - assert!(matches!( - control_rx.try_recv(), - Ok(crate::tools::TerminalSessionControl::Stop) - )); - assert!(cleanup_rx.try_recv().is_err()); - - // The terminal is only allocated after cancellation has returned. - responder.respond(CreateTerminalResponse::new("late-terminal"))?; - let mut cleanup = Vec::new(); - for _ in 0..2 { - let (method, session_id, terminal_id) = - tokio::time::timeout(Duration::from_secs(2), cleanup_rx.recv()) - .await - .expect("late terminal did not receive both cleanup requests") - .expect("mock cleanup channel closed"); - assert_eq!(session_id.to_string(), "cancelled-session"); - assert_eq!(terminal_id.to_string(), "late-terminal"); - cleanup.push(method); - } - cleanup.sort_unstable(); - assert_eq!(cleanup, ["kill", "release"]); - if !acknowledge_kill { - // Only drop the unanswered kill responder after observing release. - // Dropping it earlier would send an error response and accidentally - // allow cleanup that waits for kill acknowledgement to pass. - let pending_kill = pending_kill_rx - .try_recv() - .expect("kill must remain unanswered until release arrives"); - drop(pending_kill); - } - Ok(()) - }) - .await; - result.expect("mock ACP connection failed"); - }) - .await; - } - - #[tokio::test(flavor = "current_thread")] - async fn cancelled_terminal_creation_cleans_up_delayed_response() { - delayed_terminal_creation_cleanup(true).await; - } - - #[tokio::test(flavor = "current_thread")] - async fn cancelled_terminal_creation_releases_without_kill_acknowledgement() { - delayed_terminal_creation_cleanup(false).await; - } - - #[tokio::test] - async fn cancellation_drops_pending_generation_without_waiting() { - struct OnDrop(Arc); - impl Drop for OnDrop { - fn drop(&mut self) { - self.0.store(true, std::sync::atomic::Ordering::SeqCst); - } - } - let token = CancellationToken::new(); - let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let guard = OnDrop(dropped.clone()); - let future = async move { - let _guard = guard; - std::future::pending::().await - }; - let cancel = token.clone(); - tokio::spawn(async move { - tokio::task::yield_now().await; - cancel.cancel(); - }); - assert!(tokio::time::timeout( - std::time::Duration::from_secs(1), - cancellable(&token, future) - ) - .await - .unwrap() - .is_none()); - assert!(dropped.load(std::sync::atomic::Ordering::SeqCst)); - } - - #[tokio::test] - async fn cancellation_wins_over_already_ready_answer() { - let token = CancellationToken::new(); - token.cancel(); - assert!(cancellable(&token, std::future::ready("answer")) - .await - .is_none()); - } - - #[test] - fn previews_keep_workspace_diffs_but_do_not_read_protected_files() { - let workspace = tempfile::tempdir().unwrap(); - let external = tempfile::tempdir().unwrap(); - let ordinary = workspace.path().join("hello.txt"); - std::fs::write(&ordinary, "before").unwrap(); - let preview = - serde_json::to_string(&preflight_diff("hello.txt", "after", workspace.path())).unwrap(); - assert!(preview.contains("before")); - for path in [ - workspace.path().join(".env"), - external.path().join("secret.txt"), - ] { - std::fs::write(&path, "secret-never-preview").unwrap(); - let preview = serde_json::to_string(&preflight_diff( - path.to_str().unwrap(), - "after", - workspace.path(), - )) - .unwrap(); - assert!(!preview.contains("secret-never-preview")); - assert!(permission_edit_diff( - &serde_json::json!({ - "file_path": path, "old_string": "secret", "new_string": "public" - }), - workspace.path() - ) - .is_none()); - } - } - - #[cfg(unix)] - #[test] - fn previews_reject_symlink_escape_and_sensitive_alias() { - let workspace = tempfile::tempdir().unwrap(); - let external = tempfile::tempdir().unwrap(); - let secret = external.path().join("secret.txt"); - std::fs::write(&secret, "secret").unwrap(); - let link = workspace.path().join("alias.txt"); - std::os::unix::fs::symlink(secret, &link).unwrap(); - assert!(!preview_read_allowed(&link, workspace.path())); - let env = workspace.path().join(".env"); - std::fs::write(&env, "secret").unwrap(); - let alias = workspace.path().join("config.txt"); - std::os::unix::fs::symlink(env, &alias).unwrap(); - assert!(!preview_read_allowed(&alias, workspace.path())); - } -} diff --git a/src/acp/service/attachments.rs b/src/acp/service/attachments.rs new file mode 100644 index 0000000..4c822d3 --- /dev/null +++ b/src/acp/service/attachments.rs @@ -0,0 +1,400 @@ +use super::internal_error; +use super::AcpSession; +use crate::config::configuration::LoadedConfig; +use agent_client_protocol::schema::v1::ContentBlock; +use agent_client_protocol::schema::v1::EmbeddedResourceResource; +use agent_client_protocol::Error; +use base64::Engine as _; +use std::path::Path; + +fn write_prompt_audio( + session_id: &str, + audio: &agent_client_protocol::schema::v1::AudioContent, +) -> Result { + const MAX_AUDIO_BYTES: usize = 20 * 1024 * 1024; + let media_type = audio.mime_type.trim().to_ascii_lowercase(); + let extension = match media_type.as_str() { + "audio/wav" | "audio/x-wav" | "audio/wave" => "wav", + "audio/mpeg" | "audio/mp3" => "mp3", + _ => { + return Err(Error::invalid_params() + .data(format!("unsupported audio MIME type: {}", audio.mime_type))); + } + }; + let data = base64::engine::general_purpose::STANDARD + .decode(audio.data.trim()) + .map_err(|error| Error::invalid_params().data(format!("invalid audio data: {error}")))?; + if data.len() > MAX_AUDIO_BYTES { + return Err(Error::invalid_params().data("audio exceeds the 20 MiB size limit")); + } + let path = crate::persistence::attachments::write(session_id, extension, &data) + .map_err(|_| internal_error())?; + Ok(path.to_string_lossy().into_owned()) +} + +pub(super) fn model_supports_audio(config: &LoadedConfig, provider: &str, model: &str) -> bool { + crate::model::discovery::Discovery::new_with_custom(Some( + config.merged_config.custom_providers.clone(), + )) + .ok() + .is_some_and(|discovery| discovery.model_supports_input_modality(provider, model, "audio")) +} + +pub(super) struct ManagedAttachmentGuard { + paths: Vec, + committed: bool, +} + +impl ManagedAttachmentGuard { + pub(super) fn new(paths: Vec) -> Self { + Self { + paths, + committed: false, + } + } + + pub(super) fn commit(&mut self) { + self.committed = true; + } +} + +impl Drop for ManagedAttachmentGuard { + fn drop(&mut self) { + if !self.committed { + for path in &self.paths { + crate::persistence::attachments::remove_file(Path::new(path)); + } + } + } +} + +pub(super) fn prompt_content( + parts: Vec, + text_override: Option<&str>, + supports_images: bool, + supports_audio: bool, + session_id: &str, + session: &AcpSession, +) -> Result<(String, Vec, Vec), Error> { + let mut text = String::new(); + let mut text_override = text_override; + let mut local_image_paths = Vec::new(); + let mut local_audio_paths = Vec::new(); + for part in parts { + match part { + ContentBlock::Text(content) => { + text.push_str(text_override.take().unwrap_or(&content.text)); + } + ContentBlock::ResourceLink(link) => { + text.push_str(&format!("[{}]", link.uri)); + } + ContentBlock::Resource(resource) => match resource.resource { + EmbeddedResourceResource::TextResourceContents(resource) => { + text.push_str(&format!("[{}]\n{}", resource.uri, resource.text)); + } + EmbeddedResourceResource::BlobResourceContents(resource) => { + text.push_str(&format!("[{}]", resource.uri)); + } + _ => {} + }, + ContentBlock::Image(image) => { + if !supports_images { + return Err(Error::invalid_params().data(format!( + "model {}/{} does not support image input", + session.provider, session.model + ))); + } + match write_prompt_image(session_id, &image) { + Ok(path) => local_image_paths.push(path), + Err(error) => { + for path in &local_image_paths { + crate::persistence::attachments::remove_file(Path::new(path)); + } + return Err(error); + } + } + } + ContentBlock::Audio(audio) => { + if !supports_audio { + for path in local_image_paths.iter().chain(local_audio_paths.iter()) { + crate::persistence::attachments::remove_file(Path::new(path)); + } + return Err(Error::invalid_params().data(format!( + "model {}/{} does not support audio input", + session.provider, session.model + ))); + } + match write_prompt_audio(session_id, &audio) { + Ok(path) => local_audio_paths.push(path), + Err(error) => { + for path in local_image_paths.iter().chain(local_audio_paths.iter()) { + crate::persistence::attachments::remove_file(Path::new(path)); + } + return Err(error); + } + } + } + _ => {} + } + } + if text.is_empty() { + if !local_image_paths.is_empty() { + text.push_str("[Image attached]"); + } else if !local_audio_paths.is_empty() { + text.push_str("[Audio attached]"); + } + } + Ok((text, local_image_paths, local_audio_paths)) +} + +fn prompt_text(parts: Vec) -> Result { + let mut text = String::new(); + for part in parts { + match part { + ContentBlock::Text(content) => text.push_str(&content.text), + ContentBlock::ResourceLink(link) => text.push_str(&format!("[{}]", link.uri)), + ContentBlock::Resource(resource) => match resource.resource { + EmbeddedResourceResource::TextResourceContents(resource) => { + text.push_str(&format!("[{}]\n{}", resource.uri, resource.text)); + } + EmbeddedResourceResource::BlobResourceContents(resource) => { + text.push_str(&format!("[{}]", resource.uri)); + } + _ => {} + }, + ContentBlock::Image(_) | ContentBlock::Audio(_) => { + return Err( + Error::invalid_params().data("binary ACP prompt content is not supported here") + ); + } + _ => {} + } + } + Ok(text) +} + +fn write_prompt_image( + session_id: &str, + image: &agent_client_protocol::schema::v1::ImageContent, +) -> Result { + const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024; + + let (mime_type, encoded_data) = prompt_image_payload(image)?; + let extension = match mime_type { + "image/png" => "png", + "image/jpeg" => "jpg", + "image/gif" => "gif", + "image/webp" => "webp", + mime_type => { + return Err( + Error::invalid_params().data(format!("unsupported image MIME type: {mime_type}")) + ); + } + }; + let data = base64::engine::general_purpose::STANDARD + .decode(encoded_data) + .map_err(|error| Error::invalid_params().data(format!("invalid image data: {error}")))?; + if data.len() > MAX_IMAGE_BYTES { + return Err(Error::invalid_params().data("image exceeds the 20 MiB size limit")); + } + + let path = crate::persistence::attachments::write(session_id, extension, &data) + .map_err(|_| internal_error())?; + Ok(path.to_string_lossy().into_owned()) +} + +fn prompt_image_payload( + image: &agent_client_protocol::schema::v1::ImageContent, +) -> Result<(&str, &str), Error> { + let data = image.data.trim(); + if data.is_empty() { + return Err(Error::invalid_params().data("image data is empty")); + } + + let Some(data_uri) = data.strip_prefix("data:") else { + return Ok((image.mime_type.as_str(), data)); + }; + let (metadata, payload) = data_uri + .split_once(',') + .ok_or_else(|| Error::invalid_params().data("invalid image data URI"))?; + let mut metadata = metadata.split(';'); + let mime_type = metadata + .next() + .filter(|mime_type| !mime_type.is_empty()) + .ok_or_else(|| Error::invalid_params().data("image data URI is missing a MIME type"))?; + if !metadata.any(|value| value.eq_ignore_ascii_case("base64")) { + return Err(Error::invalid_params().data("image data URI must be base64 encoded")); + } + let payload = payload.trim(); + if payload.is_empty() { + return Err(Error::invalid_params().data("image data is empty")); + } + + Ok((mime_type, payload)) +} + +#[cfg(test)] +mod tests { + use super::super::test_support::test_session; + use super::*; + use agent_client_protocol::schema::v1::ContentBlock; + use agent_client_protocol::schema::v1::EmbeddedResourceResource; + use std::path::Path; + + #[test] + fn flattens_text_and_embedded_context() { + let text = prompt_text(vec![ + ContentBlock::from("Inspect this."), + ContentBlock::Resource(agent_client_protocol::schema::v1::EmbeddedResource::new( + EmbeddedResourceResource::TextResourceContents( + agent_client_protocol::schema::v1::TextResourceContents::new( + "fn main() {}", + "file:///tmp/main.rs", + ), + ), + )), + ]) + .expect("prompt text"); + + assert_eq!(text, "Inspect this.[file:///tmp/main.rs]\nfn main() {}"); + } + + #[test] + fn writes_supported_acp_image_to_managed_session_storage() { + let session_id = format!("acp-image-{}", cuid2::create_id()); + let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/png"); + let path = write_prompt_image(&session_id, &image).expect("image file"); + + assert!(Path::new(&path).starts_with(crate::persistence::attachments::root_dir())); + assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn command_text_override_preserves_embedded_context() { + let resource = agent_client_protocol::schema::v1::EmbeddedResource::new( + EmbeddedResourceResource::TextResourceContents( + agent_client_protocol::schema::v1::TextResourceContents::new( + "important context", + "file:///tmp/context.txt", + ), + ), + ); + let (text, images, audio) = prompt_content( + vec![ + ContentBlock::Text(agent_client_protocol::schema::v1::TextContent::new( + "/review src/lib.rs", + )), + ContentBlock::Text(agent_client_protocol::schema::v1::TextContent::new( + "\nAdditional instructions", + )), + ContentBlock::Resource(resource), + ], + Some("expanded command prompt"), + false, + false, + "test-session", + &test_session(), + ) + .expect("prompt content"); + + assert_eq!( + text, + "expanded command prompt\nAdditional instructions[file:///tmp/context.txt]\nimportant context" + ); + assert!(images.is_empty()); + assert!(audio.is_empty()); + } + + #[test] + fn writes_supported_acp_audio_to_managed_session_storage() { + let session_id = format!("acp-audio-{}", cuid2::create_id()); + let audio = agent_client_protocol::schema::v1::AudioContent::new("YXVkaW8=", "audio/wav"); + let path = write_prompt_audio(&session_id, &audio).expect("audio file"); + + assert!(path.ends_with(".wav")); + assert_eq!(std::fs::read(&path).unwrap(), b"audio"); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn rejects_audio_for_models_without_audio_modality() { + let session_id = format!("acp-audio-{}", cuid2::create_id()); + let result = prompt_content( + vec![ContentBlock::Audio( + agent_client_protocol::schema::v1::AudioContent::new("YXVkaW8=", "audio/wav"), + )], + None, + false, + false, + &session_id, + &test_session(), + ); + + assert!(result.is_err()); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn writes_acp_clipboard_image_data_uri_to_managed_storage() { + let session_id = format!("acp-image-{}", cuid2::create_id()); + let image = agent_client_protocol::schema::v1::ImageContent::new( + "data:image/png;base64,aGk=", + "application/octet-stream", + ); + let path = write_prompt_image(&session_id, &image).expect("image file"); + + assert!(path.ends_with(".png")); + assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn prompt_image_failure_rolls_back_prior_managed_files() { + let session_id = format!("acp-image-{}", cuid2::create_id()); + let result = prompt_content( + vec![ + ContentBlock::Image(agent_client_protocol::schema::v1::ImageContent::new( + "aGk=", + "image/png", + )), + ContentBlock::Image(agent_client_protocol::schema::v1::ImageContent::new( + "not-base64", + "image/png", + )), + ], + None, + true, + false, + &session_id, + &test_session(), + ); + + assert!(result.is_err()); + let directory = crate::persistence::attachments::session_dir(&session_id).unwrap(); + assert!( + !directory.exists() + || std::fs::read_dir(&directory) + .unwrap() + .all(|entry| entry.is_err()) + ); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn rejects_non_base64_acp_image_data_uri() { + let image = agent_client_protocol::schema::v1::ImageContent::new( + "data:image/png,not-base64", + "image/png", + ); + + assert!(write_prompt_image("test", &image).is_err()); + } + + #[test] + fn rejects_unsupported_acp_image_mime_type() { + let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/tiff"); + + assert!(write_prompt_image("test", &image).is_err()); + } +} diff --git a/src/acp/service/commands.rs b/src/acp/service/commands.rs new file mode 100644 index 0000000..30713f8 --- /dev/null +++ b/src/acp/service/commands.rs @@ -0,0 +1,465 @@ +use super::cancellable; +use super::internal_error; +use super::internal_error_with; +use super::notifications::send_text; +use super::AcpService; +use super::AcpSession; +use agent_client_protocol::schema::v1::AvailableCommand; +use agent_client_protocol::schema::v1::AvailableCommandInput; +use agent_client_protocol::schema::v1::AvailableCommandsUpdate; +use agent_client_protocol::schema::v1::ContentBlock; +use agent_client_protocol::schema::v1::PromptResponse; +use agent_client_protocol::schema::v1::StopReason; +use agent_client_protocol::schema::v1::UnstructuredCommandInput; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use agent_client_protocol::Error; +use tokio_util::sync::CancellationToken; + +pub(super) fn command_text(parts: &[ContentBlock]) -> String { + parts + .iter() + .find_map(|part| match part { + ContentBlock::Text(content) => Some(content.text.clone()), + _ => None, + }) + .unwrap_or_default() +} + +pub(super) fn compact_command(prompt: &str) -> Result { + let trimmed = prompt.trim(); + let Some(command_line) = trimmed.strip_prefix('/') else { + return Ok(false); + }; + let (name, args) = command_line + .split_once(char::is_whitespace) + .map(|(name, args)| (name, args.trim())) + .unwrap_or((command_line, "")); + if name != "compact" { + return Ok(false); + } + if !args.is_empty() { + return Err(Error::invalid_params().data("Usage: /compact")); + } + Ok(true) +} + +fn available_commands(session: &AcpSession) -> Vec { + let mut commands: Vec<_> = session + .config + .merged_config + .commands + .iter() + .filter(|command| !matches!(command.name.as_str(), "btw" | "compact" | "skills" | "mcp")) + .map(|command| { + let description = command + .description + .clone() + .unwrap_or_else(|| format!("Run /{}", command.name)); + AvailableCommand::new(command.name.clone(), description).input( + AvailableCommandInput::Unstructured(UnstructuredCommandInput::new("Arguments")), + ) + }) + .collect(); + commands.extend( + session + .skills + .all() + .into_iter() + .filter(|skill| !matches!(skill.name.as_str(), "btw" | "compact" | "skills" | "mcp")) + .map(|skill| { + AvailableCommand::new( + skill.name.clone(), + skill + .description + .clone() + .unwrap_or_else(|| format!("Use the {} skill", skill.name)), + ) + .input(AvailableCommandInput::Unstructured( + UnstructuredCommandInput::new("Task or context for this skill"), + )) + }), + ); + commands.push(AvailableCommand::new( + "skills", + "List skills available in this workspace", + )); + commands.push(AvailableCommand::new( + "mcp", + "List configured MCP servers and their status", + )); + commands.push(AvailableCommand::new( + "compact", + "Summarize this session to reduce context", + )); + commands.push( + AvailableCommand::new( + "btw", + "Ask a quick side question without changing session history", + ) + .input(AvailableCommandInput::Unstructured( + UnstructuredCommandInput::new("Question"), + )), + ); + commands.sort_by(|left, right| left.name.cmp(&right.name)); + commands.dedup_by(|left, right| left.name == right.name); + commands +} + +#[derive(Debug, PartialEq, Eq)] +pub(super) enum SlashExpansion { + Prompt { + prompt: String, + agent: Option, + model: Option, + }, + LocalResult(String), + Btw(String), +} + +pub(super) async fn expand_slash_command( + session: &AcpSession, + prompt: &str, +) -> Result { + let Some(command_line) = prompt.strip_prefix('/') else { + return Ok(SlashExpansion::Prompt { + prompt: prompt.to_string(), + agent: None, + model: None, + }); + }; + let (name, args) = command_line + .split_once(char::is_whitespace) + .map(|(name, args)| (name, args.trim_start())) + .unwrap_or((command_line, "")); + if name == "btw" { + let question = args.trim(); + if question.is_empty() { + return Err(Error::invalid_params().data("Usage: /btw ")); + } + return Ok(SlashExpansion::Btw(question.to_string())); + } + if name == "skills" { + if !args.is_empty() { + return Err(Error::invalid_params().data("Usage: /skills")); + } + let skills = session + .skills + .all() + .into_iter() + .map(|skill| { + format!( + "- /{} — {}", + skill.name, + skill.description.as_deref().unwrap_or("No description") + ) + }) + .collect::>(); + return Ok(SlashExpansion::LocalResult(if skills.is_empty() { + "No skills are available in this workspace.".to_string() + } else { + format!("Available skills:\n{}", skills.join("\n")) + })); + } + if name == "mcp" { + if !args.is_empty() { + return Err(Error::invalid_params().data("Usage: /mcp")); + } + let manager = crate::mcp::McpManager::ensure( + session.config.merged_config.mcp.clone(), + session.cwd.clone(), + ); + let servers = manager + .lock() + .await + .views() + .into_iter() + .map(|server| { + let detail = server + .detail + .map(|detail| format!(": {detail}")) + .unwrap_or_default(); + format!( + "- {} ({}, {}){}", + server.name, server.kind, server.status, detail + ) + }) + .collect::>(); + return Ok(SlashExpansion::LocalResult(if servers.is_empty() { + "No MCP servers are configured for this workspace.".to_string() + } else { + format!("MCP servers:\n{}", servers.join("\n")) + })); + } + if let Some(command) = session + .config + .merged_config + .commands + .iter() + .find(|command| command.name == name) + { + return command + .render(args) + .await + .map_err(|_| internal_error()) + .map(|rendered| SlashExpansion::Prompt { + prompt: rendered.prompt, + agent: rendered.agent, + model: rendered.model, + }); + } + if let Some(skill) = session.skills.get(name) { + let mut expanded = skill.content.clone(); + if !args.is_empty() { + expanded.push_str("\n\nUser task/context:\n"); + expanded.push_str(args); + } + return Ok(SlashExpansion::Prompt { + prompt: expanded, + agent: None, + model: None, + }); + } + Err(Error::invalid_params().data(format!("Unknown ACP command: /{name}"))) +} + +impl AcpService { + pub async fn available_commands( + &self, + session_id: &str, + ) -> Result { + let sessions = self.sessions.lock().await; + let session = sessions + .attachments + .get(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + Ok(AvailableCommandsUpdate::new(available_commands(session))) + } + + pub(super) async fn btw_session( + &self, + session_id: &str, + session: AcpSession, + question: String, + connection: ConnectionTo, + cancellation: CancellationToken, + ) -> Result { + let history = self + .session_manager + .lock() + .map_err(|_| internal_error())? + .get_session_ref(session_id) + .map(|stored| stored.messages.clone()) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + let answer = { + let result = cancellable( + &cancellation, + crate::llm::client::generate_btw_answer( + session.provider, + session.model, + question, + history, + ), + ) + .await; + let Some(result) = result else { + return Ok(PromptResponse::new(StopReason::Cancelled)); + }; + result.map_err(|error| internal_error_with(&error.to_string()))? + }; + // Serialize notification with close/delete, which cancel under this same lock. + let sessions = self.sessions.lock().await; + if cancellation.is_cancelled() || !sessions.attachments.contains_key(session_id) { + return Ok(PromptResponse::new(StopReason::Cancelled)); + } + send_text(&connection, session_id, &cuid2::create_id(), answer, false)?; + Ok(PromptResponse::new(StopReason::EndTurn)) + } +} + +#[cfg(test)] +mod tests { + use super::super::test_support::config_with_command; + use super::super::test_support::session_with_config; + use super::super::test_support::test_session; + use super::*; + use agent_client_protocol::schema::v1::AvailableCommandInput; + use std::path::PathBuf; + + #[test] + fn recognizes_only_exact_compact_control_command() { + assert_eq!(compact_command("/compact").unwrap(), true); + assert_eq!(compact_command(" /compact ").unwrap(), true); + assert_eq!(compact_command("/compactness").unwrap(), false); + assert_eq!(compact_command("hello").unwrap(), false); + assert!(compact_command("/compact extra").is_err()); + } + + #[test] + fn advertises_compact_as_no_input_command() { + let session = test_session(); + let command = available_commands(&session) + .into_iter() + .find(|command| command.name == "compact") + .expect("compact command"); + assert_eq!( + command.description, + "Summarize this session to reduce context" + ); + assert!(command.input.is_none()); + } + + #[tokio::test] + async fn advertises_and_parses_btw_side_command() { + let session = test_session(); + let command = available_commands(&session) + .into_iter() + .find(|command| command.name == "btw") + .expect("btw command"); + assert!(command.input.is_some()); + assert_eq!( + expand_slash_command(&session, "/btw what changed?") + .await + .expect("btw expansion"), + SlashExpansion::Btw("what changed?".to_string()) + ); + assert!(expand_slash_command(&session, "/btw").await.is_err()); + } + + #[test] + fn advertises_custom_commands_with_argument_input() { + let config = config_with_command(crate::command::custom::CustomCommand { + name: "review".to_string(), + description: Some("Review selected code".to_string()), + template: "Review: $ARGUMENTS".to_string(), + agent: None, + model: None, + subtask: Some(false), + source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( + "/tmp/opencode.jsonc", + )), + workdir: PathBuf::from("/tmp"), + }); + + let session = session_with_config(config); + let commands = available_commands(&session); + + let command = commands + .iter() + .find(|command| command.name == "review") + .expect("review command"); + assert!(commands.iter().any(|command| command.name == "skills")); + assert!(commands.iter().any(|command| command.name == "mcp")); + assert_eq!(command.description, "Review selected code"); + assert!(matches!( + command.input, + Some(AvailableCommandInput::Unstructured(_)) + )); + } + + #[tokio::test] + async fn expands_custom_slash_command_before_prompting() { + let config = config_with_command(crate::command::custom::CustomCommand { + name: "review".to_string(), + description: None, + template: "Review this carefully: $ARGUMENTS".to_string(), + agent: Some("build".to_string()), + model: Some("openai/gpt-5".to_string()), + subtask: Some(false), + source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( + "/tmp/opencode.jsonc", + )), + workdir: PathBuf::from("/tmp"), + }); + + let session = session_with_config(config); + let prompt = expand_slash_command(&session, "/review src/acp/service.rs") + .await + .expect("expanded command"); + + assert_eq!( + prompt, + SlashExpansion::Prompt { + prompt: "Review this carefully: src/acp/service.rs".to_string(), + agent: Some("build".to_string()), + model: Some("openai/gpt-5".to_string()), + } + ); + } + + #[tokio::test] + async fn advertises_and_expands_workspace_skills() { + let temp = tempfile::tempdir().expect("temp dir"); + let skill_dir = temp.path().join(".opencode/skill/reviewer"); + std::fs::create_dir_all(&skill_dir).expect("skill dir"); + std::fs::write( + skill_dir.join("SKILL.md"), + "---\nname: reviewer\ndescription: Review code carefully\n---\nInspect correctness and risks.", + ) + .expect("skill file"); + let mut config = config_with_command(crate::command::custom::CustomCommand { + name: "custom".to_string(), + description: None, + template: "$ARGUMENTS".to_string(), + agent: None, + model: None, + subtask: Some(false), + source: crate::command::custom::CustomCommandSource::Config( + temp.path().join("opencode.jsonc"), + ), + workdir: temp.path().to_path_buf(), + }); + config.project_root = temp.path().to_path_buf(); + config.cwd = temp.path().to_path_buf(); + config.xdg_config_home = temp.path().join("config"); + let session = session_with_config(config); + + let commands = available_commands(&session); + assert!(commands.iter().any(|command| command.name == "reviewer")); + let prompt = expand_slash_command(&session, "/reviewer src/lib.rs") + .await + .expect("expanded skill"); + let SlashExpansion::Prompt { prompt, .. } = prompt else { + panic!("expected skill prompt"); + }; + assert!(prompt.contains("Inspect correctness and risks.")); + assert!(prompt.contains("src/lib.rs")); + } + + #[tokio::test] + async fn reports_configured_mcp_servers() { + let mut config = config_with_command(crate::command::custom::CustomCommand { + name: "custom".to_string(), + description: None, + template: "$ARGUMENTS".to_string(), + agent: None, + model: None, + subtask: Some(false), + source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( + "/tmp/opencode.jsonc", + )), + workdir: PathBuf::from("/tmp"), + }); + config.merged_config.mcp.insert( + "filesystem".to_string(), + crate::config::configuration::McpServerConfig::Local( + crate::config::configuration::McpLocalConfig { + command: vec!["mcp-server".to_string()], + cwd: None, + environment: Default::default(), + enabled: true, + timeout_ms: None, + }, + ), + ); + let session = session_with_config(config); + + let prompt = expand_slash_command(&session, "/mcp") + .await + .expect("mcp status"); + let SlashExpansion::LocalResult(prompt) = prompt else { + panic!("expected local MCP result"); + }; + assert!(prompt.contains("filesystem (local, connecting)")); + } +} diff --git a/src/acp/service/compaction.rs b/src/acp/service/compaction.rs new file mode 100644 index 0000000..10b1cf0 --- /dev/null +++ b/src/acp/service/compaction.rs @@ -0,0 +1,610 @@ +use super::cancellable; +use super::config::model_output_limit; +use super::config::model_reasoning_capability; +use super::internal_error; +use super::internal_error_with; +use super::notifications::send_text; +use super::usage::acp_context_tokens; +use super::usage::acp_session_cost; +use super::usage::estimate_session_usage_cost; +use super::usage::send_usage; +use super::AcpService; +use super::AcpSession; +use agent_client_protocol::schema::v1::PromptResponse; +use agent_client_protocol::schema::v1::StopReason; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use agent_client_protocol::Error; +use tokio_util::sync::CancellationToken; + +fn compaction_reasoning(session: &AcpSession) -> Option { + use crate::model::reasoning::ReasoningEffort; + let capability = model_reasoning_capability( + &session.config, + &session.models, + &session.provider, + &session.model, + )?; + [ + ReasoningEffort::None, + ReasoningEffort::Minimal, + ReasoningEffort::Low, + ] + .into_iter() + .find(|effort| capability.values().contains(effort)) + .or(session.reasoning) + .filter(|effort| *effort != ReasoningEffort::None) +} + +fn compacted_messages( + messages: &[crate::session::types::Message], + selection: &crate::session::compaction::CompactionSelection, + summary: &crate::llm::client::CompactionSummary, + session: &AcpSession, + before_tokens: usize, + before_messages: usize, +) -> Result< + ( + Vec, + crate::session::types::CompactionStats, + ), + Error, +> { + let mut compacted = crate::session::compaction::apply_soft_compaction( + messages, + selection, + &summary.text, + Some(session.model.clone()), + Some(session.provider.clone()), + Some(session.agent.clone()), + crate::session::types::CompactionStats { + before_tokens, + after_tokens: 0, + before_messages, + after_messages: 0, + }, + ); + crate::session::compaction::attach_summary_usage( + &mut compacted, + crate::session::types::RecordedUsage { + input: summary.usage.input, + output: summary.usage.output, + cache_read: summary.usage.cache_read, + cache_write: summary.usage.cache_write, + cost: estimate_session_usage_cost(session, &summary.usage), + }, + ); + let after_tokens = crate::session::compaction::total_context_tokens(&compacted); + let after_messages = crate::session::compaction::filter_messages_for_context(&compacted).len(); + let stats = crate::session::types::CompactionStats { + before_tokens, + after_tokens, + before_messages, + after_messages, + }; + if after_tokens >= before_tokens { + return Err(Error::invalid_params().data(format!( + "Compaction did not reduce context ({})", + crate::session::compaction::format_compaction_stats(stats) + ))); + } + if let Some(marker) = compacted + .iter_mut() + .rev() + .find(|message| crate::session::compaction::is_compaction_marker(message)) + { + marker.compaction_stats = Some(stats); + } + Ok((compacted, stats)) +} + +impl AcpService { + pub(super) async fn compact_session( + &self, + session_id: &str, + session: AcpSession, + connection: ConnectionTo, + cancellation: CancellationToken, + ) -> Result { + let status_result = self + .session_manager + .lock() + .map_err(|_| internal_error())? + .set_session_status( + session_id, + crate::session::types::SessionStatus::Streaming, + None, + ); + if status_result.is_err() { + return Err(internal_error()); + } + + let result = self + .run_compaction(session_id, &session, cancellation.clone(), 0) + .await; + self.session_manager + .lock() + .map_err(|_| internal_error())? + .set_session_status( + session_id, + if cancellation.is_cancelled() { + crate::session::types::SessionStatus::Interrupted + } else { + crate::session::types::SessionStatus::Idle + }, + None, + ) + .map_err(|_| internal_error())?; + + match result { + Ok(stats) => { + self.notify_compaction(session_id, &session, &connection, stats, 0)?; + Ok(PromptResponse::new(StopReason::EndTurn)) + } + Err(_error) if cancellation.is_cancelled() => { + Ok(PromptResponse::new(StopReason::Cancelled)) + } + Err(error) => Err(error), + } + } + + async fn run_compaction( + &self, + session_id: &str, + session: &AcpSession, + cancellation: CancellationToken, + minimum_tokens: usize, + ) -> Result { + let messages = { + let manager = self.session_manager.lock().map_err(|_| internal_error())?; + manager + .get_session_ref(session_id) + .map(|stored| stored.messages.clone()) + .ok_or_else(|| Error::invalid_params().data("unknown session"))? + }; + let selection = crate::session::compaction::select_messages_for_compaction_with_min( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + minimum_tokens, + ) + .ok_or_else(|| Error::invalid_params().data("Nothing to compact"))?; + let before_tokens = crate::session::compaction::total_context_tokens(&messages); + let before_messages = + crate::session::compaction::filter_messages_for_context(&messages).len(); + let prompt = crate::session::compaction::build_prompt(&selection.messages_to_summarize); + let summary = cancellable( + &cancellation, + crate::llm::client::summarize_for_compaction( + session.provider.clone(), + session.model.clone(), + compaction_reasoning(session), + prompt, + cancellation.clone(), + ), + ) + .await + .ok_or_else(|| internal_error_with("Compaction cancelled by user"))? + .map_err(|error| internal_error_with(&error.to_string()))?; + if cancellation.is_cancelled() { + return Err(internal_error_with("Compaction cancelled by user")); + } + let (compacted, stats) = compacted_messages( + &messages, + &selection, + &summary, + session, + before_tokens, + before_messages, + )?; + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + manager + .replace_session_messages(session_id, compacted) + .map_err(|_| internal_error())?; + Ok(stats) + } + + pub(super) async fn maybe_auto_compact_session( + &self, + session_id: &str, + session: &AcpSession, + connection: &ConnectionTo, + cancellation: CancellationToken, + prompt_tokens: usize, + additional_tokens: usize, + ) -> Result { + let messages = self + .session_manager + .lock() + .map_err(|_| internal_error())? + .get_session_ref(session_id) + .map(|stored| stored.messages.clone()) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + let used_tokens = + acp_context_tokens(&messages, prompt_tokens).saturating_add(additional_tokens); + if !crate::session::compaction::should_auto_compact( + &session.config.merged_config.compaction, + used_tokens, + session.context_window, + model_output_limit(&session.config, &session.provider, &session.model), + ) { + return Ok(false); + } + if crate::session::compaction::select_messages_for_compaction( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + ) + .is_none() + { + return Ok(false); + } + self.session_manager + .lock() + .map_err(|_| internal_error())? + .set_session_status( + session_id, + crate::session::types::SessionStatus::Streaming, + None, + ) + .map_err(|_| internal_error())?; + let result = self + .run_compaction( + session_id, + session, + cancellation.clone(), + crate::session::compaction::MIN_COMPACTABLE_TOKENS, + ) + .await; + self.session_manager + .lock() + .map_err(|_| internal_error())? + .set_session_status( + session_id, + if cancellation.is_cancelled() { + crate::session::types::SessionStatus::Interrupted + } else { + crate::session::types::SessionStatus::Idle + }, + None, + ) + .map_err(|_| internal_error())?; + match result { + Ok(stats) => { + self.notify_compaction(session_id, session, connection, stats, prompt_tokens)?; + Ok(true) + } + Err(error) => { + crate::emit_log!("ACP auto-compaction skipped after failure: {:?}", error); + Ok(false) + } + } + } + + fn notify_compaction( + &self, + session_id: &str, + session: &AcpSession, + connection: &ConnectionTo, + stats: crate::session::types::CompactionStats, + prompt_tokens: usize, + ) -> Result<(), Error> { + let (used_tokens, cost) = { + let manager = self.session_manager.lock().map_err(|_| internal_error())?; + let stored = manager + .get_session_ref(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + ( + acp_context_tokens(&stored.messages, prompt_tokens), + acp_session_cost(&stored.messages), + ) + }; + send_text( + connection, + session_id, + &cuid2::create_id(), + format!( + "Context compacted ({})", + crate::session::compaction::format_compaction_stats(stats) + ), + false, + )?; + send_usage( + connection, + session_id, + session, + used_tokens, + (cost > 0.0).then_some(cost), + None, + ) + } +} + +#[cfg(test)] +mod tests { + use super::super::prompt::acp_request_messages; + use super::super::test_support::test_session; + use super::super::usage::acp_context_tokens; + use super::super::usage::acp_session_cost; + use super::super::AcpService; + use super::super::AcpSessions; + use super::*; + use crate::session::manager::SessionManager; + use agent_client_protocol::schema::v1::SessionNotification; + use agent_client_protocol::schema::v1::SessionUpdate; + use agent_client_protocol::Client; + use std::sync::Arc; + use std::sync::Mutex; + use tokio::sync::Mutex as AsyncMutex; + + #[test] + fn acp_compaction_discards_provider_usage_before_compaction_boundary() { + use crate::session::types::{CompactionStats, Message, MessagePart}; + let mut assistant = Message::assistant("old answer"); + assistant + .parts + .push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); + let mut summary = Message::user(format!( + "{}\nsummary", + crate::session::compaction::SUMMARY_PREFIX + )); + summary + .parts + .push(MessagePart::usage(95_000, 500, 0, 0, 0.0)); + let marker = crate::session::compaction::compaction_marker(CompactionStats { + before_tokens: 90_100, + after_tokens: 100, + before_messages: 2, + after_messages: 1, + }); + let messages = vec![Message::user("question"), assistant, summary, marker]; + + assert_eq!( + acp_context_tokens(&messages, 1_000), + crate::session::compaction::total_context_tokens(&messages) + 1_000 + ); + } + + #[test] + fn acp_compaction_honors_reserved_tokens_and_auto_setting() { + use crate::config::configuration::CompactionConfig; + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(70_000, 0, 10_000, 0, 0.0)); + let messages = vec![Message::user("history ".repeat(2_000)), assistant]; + assert!(crate::session::compaction::select_messages_for_compaction( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + ) + .is_some()); + let config = CompactionConfig::Settings { + auto: true, + prune: false, + reserved: Some(20_000), + }; + let used = acp_context_tokens(&messages, 1_000); + assert!(crate::session::compaction::total_context_tokens(&messages) < 80_000); + assert!(crate::session::compaction::should_auto_compact( + &config, + used, + Some(100_000), + None + )); + assert!(!crate::session::compaction::should_auto_compact( + &config, + used, + Some(100_001), + None + )); + assert!(!crate::session::compaction::should_auto_compact( + &CompactionConfig::Disabled, + used, + Some(100_000), + None + )); + assert!(!crate::session::compaction::should_auto_compact( + &config, used, None, None + )); + let no_auto = CompactionConfig::Settings { + auto: false, + prune: false, + reserved: Some(20_000), + }; + assert!(!crate::session::compaction::should_auto_compact( + &no_auto, + used, + Some(100_000), + None + )); + } + + #[test] + fn acp_compaction_keeps_fresh_system_prompt_in_request() { + use crate::session::types::{CompactionStats, Message}; + let summary = Message::user(format!( + "{}\nsummary", + crate::session::compaction::SUMMARY_PREFIX + )); + let tail = Message::assistant("retained answer"); + let marker = crate::session::compaction::compaction_marker(CompactionStats { + before_tokens: 10_000, + after_tokens: 100, + before_messages: 3, + after_messages: 2, + }); + let pending = Message::user("continue"); + let history = vec![ + Message::user("archived history"), + summary.clone(), + tail.clone(), + marker, + pending.clone(), + ]; + let system = Message::system("fresh system instructions"); + let request = acp_request_messages(&history, system.clone()); + assert_eq!(request, vec![system, summary, tail, pending]); + assert_eq!( + crate::session::compaction::filter_messages_for_context(&request), + request + ); + } + + #[test] + fn acp_compaction_cost_keeps_summary_billing_without_double_counting() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant.cost = Some(0.25); + assistant + .parts + .push(MessagePart::usage(100, 20, 0, 0, 0.25)); + let mut summary = Message::user(format!( + "{}\nsummary", + crate::session::compaction::SUMMARY_PREFIX + )); + summary.parts.push(MessagePart::usage(100, 20, 0, 0, 0.05)); + assert!((acp_session_cost(&[assistant, summary]) - 0.30).abs() < f64::EPSILON); + } + + #[tokio::test(flavor = "current_thread")] + async fn acp_compaction_notifies_client_of_reduced_context_and_summary_cost() { + use crate::session::types::{Message, MessagePart}; + use agent_client_protocol::Agent; + tokio::task::LocalSet::new() + .run_until(async { + let workspace = tempfile::tempdir().unwrap(); + let mut manager = SessionManager::new() + .with_history_for_workspace(workspace.path()) + .unwrap(); + let session_id = manager.create_session(Some("ACP compaction regression".into())); + let mut assistant = Message::assistant("old answer"); + assistant + .parts + .push(MessagePart::usage(80_000, 200, 0, 0, 0.25)); + let messages = vec![Message::user("history ".repeat(2_000)), assistant]; + let selection = crate::session::compaction::select_messages_for_compaction( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + ) + .unwrap(); + let mut session = test_session(); + session.context_window = Some(100_000); + let summary = crate::llm::client::CompactionSummary { + text: "Retained task details".into(), + usage: crate::aisdk::chunk::TokenUsage { + input: 3_000, + output: 20, + cache_read: 0, + cache_write: 0, + }, + }; + let (mut compacted, stats) = compacted_messages( + &messages, + &selection, + &summary, + &session, + crate::session::compaction::total_context_tokens(&messages), + messages.len(), + ) + .unwrap(); + let summary_message = compacted + .iter_mut() + .find(|message| crate::session::compaction::is_compaction_summary(message)) + .unwrap(); + summary_message + .parts + .push(MessagePart::usage(0, 1, 0, 0, 0.05)); + let expected_used = + crate::session::compaction::total_context_tokens(&compacted) + 1_000; + manager + .replace_session_messages(&session_id, compacted) + .unwrap(); + let service = AcpService { + sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), + session_manager: Arc::new(Mutex::new(manager)), + client_capabilities: Arc::new(Mutex::new(Default::default())), + }; + let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); + let client = Client.builder().on_receive_notification( + async move |notification: SessionNotification, _connection| { + let _ = observed_tx.send(notification.update); + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ); + Agent + .builder() + .connect_with(client, async move |connection| { + service + .notify_compaction(&session_id, &session, &connection, stats, 1_000) + .unwrap(); + let feedback = tokio::time::timeout( + std::time::Duration::from_secs(2), + observed_rx.recv(), + ) + .await + .unwrap() + .unwrap(); + assert!(matches!(feedback, SessionUpdate::AgentMessageChunk(_))); + let update = tokio::time::timeout( + std::time::Duration::from_secs(2), + observed_rx.recv(), + ) + .await + .unwrap() + .unwrap(); + let SessionUpdate::UsageUpdate(usage) = update else { + panic!("missing usage update") + }; + assert_eq!(usage.used, expected_used as u64); + assert_eq!(usage.size, 100_000); + assert!((usage.cost.unwrap().amount - 0.30).abs() < f64::EPSILON); + service + .session_manager + .lock() + .unwrap() + .try_delete_session(&session_id) + .unwrap(); + Ok(()) + }) + .await + .unwrap(); + }) + .await; + } + + #[test] + fn builds_smaller_soft_compaction_for_acp() { + let session = test_session(); + let messages = vec![ + crate::session::types::Message::user("u".repeat(8_000)), + crate::session::types::Message::assistant("a".repeat(8_000)), + crate::session::types::Message::user("recent"), + ]; + let selection = crate::session::compaction::select_messages_for_compaction_with_min( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + 0, + ) + .expect("compaction selection"); + let before_tokens = crate::session::compaction::total_context_tokens(&messages); + let before_messages = + crate::session::compaction::filter_messages_for_context(&messages).len(); + + let (compacted, stats) = compacted_messages( + &messages, + &selection, + &crate::llm::client::CompactionSummary { + text: "short handoff".to_string(), + usage: Default::default(), + }, + &session, + before_tokens, + before_messages, + ) + .expect("smaller compaction"); + + assert!(stats.after_tokens < stats.before_tokens); + assert!(crate::session::compaction::latest_compaction_stats(&compacted).is_some()); + assert_eq!(compacted[0].id, messages[0].id); + } +} diff --git a/src/acp/service/config.rs b/src/acp/service/config.rs new file mode 100644 index 0000000..0cd976e --- /dev/null +++ b/src/acp/service/config.rs @@ -0,0 +1,600 @@ +use super::AcpService; +use super::AcpSession; +use crate::config::configuration::LoadedConfig; +use agent_client_protocol::schema::v1::McpServer; +use agent_client_protocol::schema::v1::SessionConfigOption; +use agent_client_protocol::schema::v1::SessionConfigOptionCategory; +use agent_client_protocol::schema::v1::SessionConfigSelectGroup; +use agent_client_protocol::schema::v1::SessionConfigSelectOption; +use agent_client_protocol::schema::v1::SessionMode; +use agent_client_protocol::schema::v1::SessionModeState; +use agent_client_protocol::schema::v1::SetSessionConfigOptionResponse; +use agent_client_protocol::Error; + +pub(super) fn model_output_limit( + config: &LoadedConfig, + provider_id: &str, + model_id: &str, +) -> Option { + if let Some(limit) = config + .merged_config + .custom_providers + .get(&provider_id.trim().to_ascii_lowercase()) + .and_then(|provider| provider.models.get(model_id)) + .and_then(|model| model.max_tokens) + { + return Some(limit); + } + crate::model::discovery::Discovery::new_with_config( + Some(config.merged_config.custom_providers.clone()), + config.merged_config.disabled_providers.clone(), + config.merged_config.enabled_providers.clone(), + ) + .ok() + .and_then(|discovery| discovery.get_model_output_limit(provider_id, model_id)) +} + +pub(super) fn resolved_reasoning( + session: &AcpSession, + requested: crate::model::reasoning::ReasoningEffort, +) -> Option { + model_reasoning_capability( + &session.config, + &session.models, + &session.provider, + &session.model, + ) + .and_then(|capability| capability.resolve(Some(requested))) +} + +fn reasoning_effort_label(effort: crate::model::reasoning::ReasoningEffort) -> &'static str { + match effort { + crate::model::reasoning::ReasoningEffort::None => "None", + crate::model::reasoning::ReasoningEffort::Minimal => "Minimal", + crate::model::reasoning::ReasoningEffort::Low => "Low", + crate::model::reasoning::ReasoningEffort::Medium => "Medium", + crate::model::reasoning::ReasoningEffort::High => "High", + crate::model::reasoning::ReasoningEffort::XHigh => "Xhigh", + crate::model::reasoning::ReasoningEffort::Max => "Max", + } +} + +pub(super) fn merge_acp_mcp_servers(config: &mut LoadedConfig, servers: Vec) { + for server in servers { + let (name, server) = match server { + McpServer::Stdio(server) => { + let mut command = vec![server.command.to_string_lossy().into_owned()]; + command.extend(server.args); + ( + server.name, + crate::config::configuration::McpServerConfig::Local( + crate::config::configuration::McpLocalConfig { + command, + cwd: None, + environment: server + .env + .into_iter() + .map(|variable| (variable.name, variable.value)) + .collect(), + enabled: true, + timeout_ms: None, + }, + ), + ) + } + McpServer::Http(server) => ( + server.name, + crate::config::configuration::McpServerConfig::Remote( + crate::config::configuration::McpRemoteConfig { + url: server.url, + headers: server + .headers + .into_iter() + .map(|header| (header.name, header.value)) + .collect(), + enabled: true, + timeout_ms: None, + oauth_enabled: false, + oauth_client_id: None, + oauth_client_secret: None, + oauth_scope: None, + }, + ), + ), + McpServer::Sse(server) => ( + server.name, + crate::config::configuration::McpServerConfig::Remote( + crate::config::configuration::McpRemoteConfig { + url: server.url, + headers: server + .headers + .into_iter() + .map(|header| (header.name, header.value)) + .collect(), + enabled: true, + timeout_ms: None, + oauth_enabled: false, + oauth_client_id: None, + oauth_client_secret: None, + oauth_scope: None, + }, + ), + ), + #[allow(unreachable_patterns)] + _ => continue, + }; + config.merged_config.mcp.insert(name, server); + } +} + +pub(super) fn resolve_model(config: &LoadedConfig) -> (String, String) { + if let Ok(Some((provider, model))) = + crate::persistence::PrefsDAO::new().and_then(|prefs| prefs.get_active_model()) + { + return (provider, model); + } + config + .merged_config + .model + .as_deref() + .map(crate::app::parse_model_ref) + .unwrap_or_else(|| ("opencode".to_string(), "big-pickle".to_string())) +} + +fn mode_description(agent: &crate::agent::definition::AgentDefinition) -> &str { + if agent.name == "compaction" && agent.description.trim().is_empty() { + "Compaction agent. Summarizes conversation context for continued work." + } else { + &agent.description + } +} + +pub(super) fn session_modes(session: &AcpSession) -> SessionModeState { + let mut modes = session + .config + .merged_config + .agent_registry + .visible_primary_agents() + .into_iter() + .map(|agent| { + SessionMode::new(agent.name.clone(), agent.name.clone()) + .description(mode_description(agent)) + }) + .collect::>(); + modes.sort_by(|left, right| left.name.cmp(&right.name)); + SessionModeState::new(session.agent.clone(), modes) +} + +pub(super) fn session_config_options(session: &AcpSession) -> Vec { + let mut mode_options = session + .config + .merged_config + .agent_registry + .visible_primary_agents() + .into_iter() + .map(|agent| { + SessionConfigSelectOption::new(agent.name.clone(), agent.name.clone()) + .description(mode_description(agent)) + }) + .collect::>(); + mode_options.sort_by(|left, right| left.name.cmp(&right.name)); + + let mut options = vec![ + SessionConfigOption::select("mode", "Mode", session.agent.clone(), mode_options) + .category(SessionConfigOptionCategory::Mode), + model_config_option(&session.models, &session.provider, &session.model), + ]; + options.push(reasoning_config_option(session)); + options +} + +fn reasoning_config_option(session: &AcpSession) -> SessionConfigOption { + let options = [ + crate::model::reasoning::ReasoningEffort::None, + crate::model::reasoning::ReasoningEffort::Low, + crate::model::reasoning::ReasoningEffort::Medium, + crate::model::reasoning::ReasoningEffort::High, + crate::model::reasoning::ReasoningEffort::XHigh, + crate::model::reasoning::ReasoningEffort::Max, + ] + .iter() + .map(|effort| SessionConfigSelectOption::new(effort.as_str(), reasoning_effort_label(*effort))) + .collect::>(); + let current = session.reasoning_selection.as_str(); + SessionConfigOption::select("effort", "Effort", current, options) + .category(SessionConfigOptionCategory::ThoughtLevel) +} + +fn model_config_option( + models: &[crate::model::types::Model], + provider: &str, + model_id: &str, +) -> SessionConfigOption { + let mut model_groups: Vec = Vec::new(); + for model in models { + let option = SessionConfigSelectOption::new(model_value(model), model.name.clone()) + .description(model.id.clone()); + if model_groups + .last() + .is_some_and(|group| group.group.to_string() == model.provider_id) + { + model_groups + .last_mut() + .expect("model group exists") + .options + .push(option); + } else { + model_groups.push(SessionConfigSelectGroup::new( + model.provider_id.clone(), + model.provider_name.clone(), + vec![option], + )); + } + } + + SessionConfigOption::select( + "model", + "Model", + format!("{provider}/{model_id}"), + model_groups, + ) + .category(SessionConfigOptionCategory::Model) +} + +fn model_value(model: &crate::model::types::Model) -> String { + crate::model::catalog::model_ref(model) +} + +pub(super) fn find_selectable_model<'a>( + models: &'a [crate::model::types::Model], + model_ref: &str, +) -> Result<&'a crate::model::types::Model, Error> { + models + .iter() + .find(|model| model_value(model) == model_ref) + .ok_or_else(|| Error::invalid_params().data("unknown or unavailable ACP model")) +} + +pub(super) fn model_reasoning( + config: &LoadedConfig, + models: &[crate::model::types::Model], + provider: &str, + model_id: &str, +) -> Option { + model_reasoning_capability(config, models, provider, model_id) + .and_then(|capability| capability.resolve(None)) +} + +pub(super) fn model_reasoning_capability( + config: &LoadedConfig, + models: &[crate::model::types::Model], + provider: &str, + model_id: &str, +) -> Option { + if let Some(capability) = models + .iter() + .find(|model| model.provider_id == provider && model.id == model_id) + .and_then(|model| { + crate::model::reasoning::capability_from_options(&model.reasoning_options) + }) + { + return Some(capability); + } + crate::model::discovery::Discovery::new_with_custom(Some( + config.merged_config.custom_providers.clone(), + )) + .ok() + .and_then(|discovery| discovery.get_model_reasoning_capability(provider, model_id)) + .filter(|capability| !capability.values().is_empty()) +} + +pub(super) fn model_context_window( + config: &LoadedConfig, + models: &[crate::model::types::Model], + provider: &str, + model: &str, +) -> Option { + if let Some(context_window) = models + .iter() + .find(|candidate| candidate.provider_id == provider && candidate.id == model) + .and_then(|model| model.context_window) + { + return Some(context_window); + } + + let discovery = crate::model::discovery::Discovery::new_with_custom(Some( + config.merged_config.custom_providers.clone(), + )) + .ok(); + discovery + .as_ref() + .and_then(|discovery| discovery.get_model_limit(provider, model)) + .or_else(|| { + config + .merged_config + .custom_providers + .get(provider) + .and_then(|provider| provider.models.get(model)) + .and_then(|model| model.context_window) + }) +} + +impl AcpService { + pub async fn set_mode( + &self, + session_id: &str, + mode_id: &str, + ) -> Result { + let mut sessions = self.sessions.lock().await; + let session = sessions + .attachments + .get_mut(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + if session + .config + .merged_config + .agent_registry + .primary_agent(mode_id) + .is_none_or(|agent| agent.hidden) + { + return Err(Error::invalid_params().data("unknown ACP mode")); + } + session.agent = mode_id.to_string(); + Ok(SetSessionConfigOptionResponse::new(session_config_options( + session, + ))) + } + + pub async fn set_model( + &self, + session_id: &str, + model_ref: &str, + ) -> Result { + let mut sessions = self.sessions.lock().await; + let session = sessions + .attachments + .get_mut(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + let model = find_selectable_model(&session.models, model_ref)?; + session.provider.clone_from(&model.provider_id); + session.model.clone_from(&model.id); + session.reasoning = resolved_reasoning(session, session.reasoning_selection); + session.context_window = model_context_window( + &session.config, + &session.models, + &session.provider, + &session.model, + ); + Ok(SetSessionConfigOptionResponse::new(session_config_options( + session, + ))) + } + + pub async fn set_reasoning_effort( + &self, + session_id: &str, + value: &str, + ) -> Result { + let mut sessions = self.sessions.lock().await; + let session = sessions + .attachments + .get_mut(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + let requested = value + .parse::() + .map_err(|_| Error::invalid_params().data("invalid reasoning effort"))?; + session.reasoning_selection = requested; + session.reasoning = resolved_reasoning(session, requested); + Ok(SetSessionConfigOptionResponse::new(session_config_options( + session, + ))) + } +} + +#[cfg(test)] +mod tests { + use super::super::test_support::config_with_command; + use super::super::test_support::empty_config; + use super::super::test_support::model; + use super::super::test_support::test_session; + use super::*; + use agent_client_protocol::schema::v1::SessionConfigOptionCategory; + use std::path::PathBuf; + + #[test] + fn merges_client_mcp_servers_into_session_config() { + let mut config = config_with_command(crate::command::custom::CustomCommand { + name: "custom".to_string(), + description: None, + template: "$ARGUMENTS".to_string(), + agent: None, + model: None, + subtask: Some(false), + source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( + "/tmp/opencode.jsonc", + )), + workdir: PathBuf::from("/tmp"), + }); + merge_acp_mcp_servers( + &mut config, + vec![serde_json::from_value(serde_json::json!({ + "name": "zed-fs", + "command": "mcp-server", + "args": ["--stdio"], + "env": [{"name": "ROOT", "value": "/workspace"}] + })) + .expect("ACP stdio MCP server")], + ); + + let server = config + .merged_config + .mcp + .get("zed-fs") + .expect("merged ACP MCP server"); + assert_eq!(server.kind(), "local"); + assert!(server.enabled()); + } + + #[test] + fn builds_grouped_model_config_option() { + let option = model_config_option( + &[ + model("anthropic", "Anthropic", "claude", "Claude"), + model("openai", "OpenAI", "gpt-5", "GPT-5"), + model("openai", "OpenAI", "gpt-5-mini", "GPT-5 Mini"), + ], + "openai", + "gpt-5", + ); + + assert_eq!(option.id.to_string(), "model"); + assert_eq!(option.category, Some(SessionConfigOptionCategory::Model)); + let agent_client_protocol::schema::v1::SessionConfigKind::Select(select) = option.kind + else { + panic!("model option should be a select"); + }; + assert_eq!(select.current_value.to_string(), "openai/gpt-5"); + let agent_client_protocol::schema::v1::SessionConfigSelectOptions::Grouped(groups) = + select.options + else { + panic!("model options should be grouped"); + }; + assert_eq!(groups.len(), 2); + assert_eq!(groups[1].name, "OpenAI"); + assert_eq!(groups[1].options[1].value.to_string(), "openai/gpt-5-mini"); + } + + #[test] + fn validates_model_refs_against_selectable_catalog() { + let models = [model("openai", "OpenAI", "gpt-5", "GPT-5")]; + + assert_eq!( + find_selectable_model(&models, "openai/gpt-5") + .expect("known model") + .id, + "gpt-5" + ); + assert!(find_selectable_model(&models, "openai/missing").is_err()); + assert!(find_selectable_model(&models, "other/gpt-5").is_err()); + } + + #[test] + fn resolves_context_window_from_selectable_models() { + let mut model = model("example", "Example", "large-context", "Large Context"); + model.context_window = Some(1_090_000); + + assert_eq!( + model_context_window(&empty_config(), &[model], "example", "large-context"), + Some(1_090_000) + ); + } + + #[test] + fn preserves_selected_reasoning_effort_when_model_cannot_apply_it() { + let session = test_session(); + let option = reasoning_config_option(&session); + assert_eq!(option.id.to_string(), "effort"); + assert_eq!(option.name, "Effort"); + assert_eq!( + option.category, + Some(SessionConfigOptionCategory::ThoughtLevel) + ); + let agent_client_protocol::schema::v1::SessionConfigKind::Select(select) = option.kind + else { + panic!("reasoning option should be a select"); + }; + assert_eq!(select.current_value.to_string(), "high"); + let agent_client_protocol::schema::v1::SessionConfigSelectOptions::Ungrouped(options) = + select.options + else { + panic!("reasoning options should be ungrouped"); + }; + let values = options + .iter() + .map(|option| (option.value.to_string(), option.name.clone())) + .collect::>(); + assert_eq!( + values, + vec![ + ("none".to_string(), "None".to_string()), + ("low".to_string(), "Low".to_string()), + ("medium".to_string(), "Medium".to_string()), + ("high".to_string(), "High".to_string()), + ("xhigh".to_string(), "Xhigh".to_string()), + ("max".to_string(), "Max".to_string()), + ] + ); + } + + #[test] + fn compaction_mode_descriptions_preserve_configuration_and_visibility() { + let mut session = test_session(); + assert!(!session_modes(&session) + .available_modes + .iter() + .any(|mode| mode.id.to_string() == "compaction")); + + for description in ["", " \n ", "Use this agent to preserve project decisions."] { + let definitions = crate::agent::definition::parse_agent_definitions_from_config( + Some(&serde_json::json!({"compaction": { + "mode": "primary", "description": description + }})), + &mut Vec::new(), + ); + session.config.merged_config.agent_registry = + crate::agent::definition::AgentRegistry::with_definitions(None, definitions); + let expected = if description.trim().is_empty() { + "Compaction agent. Summarizes conversation context for continued work." + } else { + description + }; + let modes = session_modes(&session); + let compaction = modes + .available_modes + .iter() + .find(|mode| mode.id.to_string() == "compaction") + .expect("configured compaction mode"); + assert_eq!(compaction.description.as_deref(), Some(expected)); + let config = serde_json::to_value(session_config_options(&session)).unwrap(); + let compaction = config[0]["options"] + .as_array() + .unwrap() + .iter() + .find(|option| option["value"] == "compaction") + .expect("configured compaction option"); + assert_eq!(compaction["description"], expected); + assert_eq!( + modes + .available_modes + .iter() + .find(|mode| mode.id.to_string() == "build") + .unwrap() + .description + .as_deref(), + Some("The default agent. Executes tools based on configured permissions.") + ); + } + + for settings in [ + serde_json::json!({"mode": "subagent"}), + serde_json::json!({"mode": "primary", "hidden": true}), + ] { + let definitions = crate::agent::definition::parse_agent_definitions_from_config( + Some(&serde_json::json!({"compaction": settings})), + &mut Vec::new(), + ); + session.config.merged_config.agent_registry = + crate::agent::definition::AgentRegistry::with_definitions(None, definitions); + assert!(!session_modes(&session) + .available_modes + .iter() + .any(|mode| mode.id.to_string() == "compaction")); + let config = serde_json::to_value(session_config_options(&session)).unwrap(); + assert!(!config[0]["options"] + .as_array() + .unwrap() + .iter() + .any(|option| option["value"] == "compaction")); + } + } +} diff --git a/src/acp/service/mod.rs b/src/acp/service/mod.rs new file mode 100644 index 0000000..f121f7a --- /dev/null +++ b/src/acp/service/mod.rs @@ -0,0 +1,190 @@ +mod attachments; +mod commands; +mod compaction; +mod config; +mod notifications; +mod permissions; +mod prompt; +mod questions; +mod sessions; +mod stream; +mod terminals; +mod usage; + +#[cfg(test)] +mod test_support; + +use crate::config::configuration::LoadedConfig; +use crate::session::manager::SessionManager; +use agent_client_protocol::Error; +use std::collections::HashMap; +use std::collections::HashSet; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; +use tokio::sync::Mutex as AsyncMutex; +use tokio_util::sync::CancellationToken; + +#[derive(Clone)] +pub struct AcpService { + sessions: Arc>, + session_manager: Arc>, + client_capabilities: Arc>, +} + +// Attachments detach immediately on close, but a turn retains ownership until all +// transcript/status finalization finishes. Both are protected by the same lock. +#[derive(Default)] +struct AcpSessions { + attachments: HashMap, + active_prompts: HashSet, +} + +impl AcpSessions { + fn ensure_attachable(&self, session_id: &str) -> Result<(), Error> { + if self.active_prompts.contains(session_id) { + return Err(Error::invalid_params().data("session already has an active prompt")); + } + Ok(()) + } + + fn begin_prompt(&mut self, session_id: &str) -> Result { + self.ensure_attachable(session_id)?; + let session = self + .attachments + .get_mut(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + let cancellation = CancellationToken::new(); + session.cancellation = Some(cancellation.clone()); + self.active_prompts.insert(session_id.to_string()); + Ok(cancellation) + } + + // Call only after the stream and wrapper have finished every persisted write. + fn finish_prompt(&mut self, session_id: &str) { + if let Some(session) = self.attachments.get_mut(session_id) { + session.cancellation = None; + } + self.active_prompts.remove(session_id); + } +} + +async fn cancellable( + token: &CancellationToken, + future: impl std::future::Future, +) -> Option { + tokio::select! { + biased; + _ = token.cancelled() => None, + result = future => Some(result), + } +} + +#[derive(Clone)] +struct AcpSession { + cwd: PathBuf, + config: LoadedConfig, + tool_permissions: crate::tools::ToolPermissions, + skills: crate::skill::SkillStore, + models: Vec, + provider: String, + model: String, + agent: String, + reasoning_selection: crate::model::reasoning::ReasoningEffort, + reasoning: Option, + context_window: Option, + cancellation: Option, +} + +fn internal_error() -> Error { + Error::internal_error().data("Crabcode ACP operation failed") +} + +fn internal_error_with(error: &str) -> Error { + Error::internal_error().data(format!("Crabcode ACP operation failed: {error}")) +} + +impl AcpService { + pub fn new(initial_workspace: &Path) -> Result { + let session_manager = SessionManager::new() + .with_history_for_workspace(initial_workspace) + .map_err(|_| internal_error())?; + Ok(Self { + sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), + session_manager: Arc::new(Mutex::new(session_manager)), + client_capabilities: Arc::new(Mutex::new(Default::default())), + }) + } + + pub fn set_client_capabilities( + &self, + capabilities: agent_client_protocol::schema::v1::ClientCapabilities, + ) { + if let Ok(mut current) = self.client_capabilities.lock() { + *current = capabilities; + } + } + + fn supports_form_elicitation(&self) -> bool { + self.client_capabilities + .lock() + .ok() + .and_then(|capabilities| capabilities.elicitation.clone()) + .and_then(|elicitation| elicitation.form) + .is_some() + } + + fn supports_terminals(&self) -> bool { + self.client_capabilities + .lock() + .ok() + .is_some_and(|capabilities| capabilities.terminal) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use tokio_util::sync::CancellationToken; + + #[tokio::test] + async fn cancellation_drops_pending_generation_without_waiting() { + struct OnDrop(Arc); + impl Drop for OnDrop { + fn drop(&mut self) { + self.0.store(true, std::sync::atomic::Ordering::SeqCst); + } + } + let token = CancellationToken::new(); + let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let guard = OnDrop(dropped.clone()); + let future = async move { + let _guard = guard; + std::future::pending::().await + }; + let cancel = token.clone(); + tokio::spawn(async move { + tokio::task::yield_now().await; + cancel.cancel(); + }); + assert!(tokio::time::timeout( + std::time::Duration::from_secs(1), + cancellable(&token, future) + ) + .await + .unwrap() + .is_none()); + assert!(dropped.load(std::sync::atomic::Ordering::SeqCst)); + } + + #[tokio::test] + async fn cancellation_wins_over_already_ready_answer() { + let token = CancellationToken::new(); + token.cancel(); + assert!(cancellable(&token, std::future::ready("answer")) + .await + .is_none()); + } +} diff --git a/src/acp/service/notifications.rs b/src/acp/service/notifications.rs new file mode 100644 index 0000000..ef4ec5d --- /dev/null +++ b/src/acp/service/notifications.rs @@ -0,0 +1,622 @@ +use super::internal_error; +use agent_client_protocol::schema::v1::ContentBlock; +use agent_client_protocol::schema::v1::ContentChunk; +use agent_client_protocol::schema::v1::SessionNotification; +use agent_client_protocol::schema::v1::SessionUpdate; +use agent_client_protocol::schema::v1::ToolCall; +use agent_client_protocol::schema::v1::ToolCallContent; +use agent_client_protocol::schema::v1::ToolCallLocation; +use agent_client_protocol::schema::v1::ToolCallStatus; +use agent_client_protocol::schema::v1::ToolCallUpdate; +use agent_client_protocol::schema::v1::ToolCallUpdateFields; +use agent_client_protocol::schema::v1::ToolKind; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use agent_client_protocol::Error; +use std::path::Path; +use std::path::PathBuf; + +fn mcp_native_tool_content(metadata: &serde_json::Value) -> Vec { + metadata + .get("mcp_result") + .and_then(|result| result.get("content")) + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter(|block| { + !matches!( + block.get("type").and_then(serde_json::Value::as_str), + Some("text") + ) + }) + .map(|block| { + serde_json::from_value::(block.clone()) + .map(ToolCallContent::from) + .unwrap_or_else(|_| { + let kind = block + .get("type") + .and_then(serde_json::Value::as_str) + .unwrap_or("unknown"); + ToolCallContent::from(format!( + "[MCP {kind} content]\n{}", + serde_json::to_string_pretty(block).unwrap_or_else(|_| block.to_string()) + )) + }) + }) + .collect() +} + +fn metadata_has_mcp_images(metadata: Option<&serde_json::Value>) -> bool { + metadata + .and_then(|metadata| metadata.get("mcp_result")) + .and_then(|result| result.get("content")) + .and_then(serde_json::Value::as_array) + .is_some_and(|content| { + content + .iter() + .any(|block| block.get("type").and_then(serde_json::Value::as_str) == Some("image")) + }) +} + +pub(super) fn send_text( + connection: &ConnectionTo, + session_id: &str, + message_id: &str, + text: String, + thought: bool, +) -> Result<(), Error> { + let update = if thought { + SessionUpdate::AgentThoughtChunk(ContentChunk::new(text.into()).message_id(message_id)) + } else { + SessionUpdate::AgentMessageChunk(ContentChunk::new(text.into()).message_id(message_id)) + }; + connection + .send_notification(SessionNotification::new(session_id.to_string(), update)) + .map_err(|_| internal_error()) +} + +fn tool_result_text(payload: &serde_json::Value) -> String { + payload + .get("output") + .or_else(|| payload.get("output_preview")) + .or_else(|| payload.get("error")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string() +} + +pub(super) fn absolute_tool_path(path: &str, cwd: &Path) -> PathBuf { + let path = PathBuf::from(path); + if path.is_absolute() { + path + } else { + cwd.join(path) + } +} + +pub(super) fn tool_locations( + tool_name: &str, + input: &serde_json::Value, + cwd: &Path, +) -> Vec { + let paths = if tool_name == "write_files" { + input + .get("files") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|file| file.get("file_path").and_then(serde_json::Value::as_str)) + .map(str::to_string) + .collect() + } else if tool_name == "apply_patch" { + crate::tools::patch::patch_paths_from_params(input) + } else { + input + .get("file_path") + .or_else(|| input.get("filePath")) + .or_else(|| input.get("filepath")) + .or_else(|| input.get("path")) + .and_then(serde_json::Value::as_str) + .map(|path| vec![path.to_string()]) + .unwrap_or_default() + }; + + paths + .into_iter() + .map(|path| ToolCallLocation::new(absolute_tool_path(&path, cwd))) + .collect() +} + +fn tool_result_locations(payload: &serde_json::Value, cwd: &Path) -> Vec { + let Some(metadata) = payload.get("metadata") else { + return Vec::new(); + }; + let line = metadata + .get("line_number") + .and_then(serde_json::Value::as_u64) + .and_then(|line| u32::try_from(line).ok()); + + if let Some(changes) = metadata + .get("changes") + .and_then(serde_json::Value::as_array) + { + return changes + .iter() + .filter_map(|change| change.get("path").and_then(serde_json::Value::as_str)) + .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd))) + .collect(); + } + + metadata + .get("path") + .and_then(serde_json::Value::as_str) + .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd)).line(line)) + .into_iter() + .collect() +} + +fn tool_result_content(payload: &serde_json::Value, cwd: &Path) -> Vec { + let mut content = Vec::new(); + let text = tool_result_text(payload); + if !text.is_empty() { + content.push(ToolCallContent::from(text)); + } + + if let Some(metadata) = payload.get("metadata") { + if let Some(changes) = metadata + .get("changes") + .and_then(serde_json::Value::as_array) + { + content.extend(changes.iter().filter_map(|change| tool_diff(change, cwd))); + } else if let Some(diff) = tool_diff(metadata, cwd) { + content.push(diff); + } + content.extend(mcp_native_tool_content(metadata)); + } + + if !metadata_has_mcp_images(payload.get("metadata")) { + if let Some(images) = payload.get("images").and_then(serde_json::Value::as_array) { + content.extend(images.iter().filter_map(tool_result_image)); + } + } + + content +} + +fn tool_diff(change: &serde_json::Value, cwd: &Path) -> Option { + let path = change.get("path")?.as_str()?; + let new_text = change.get("new_text")?.as_str()?; + let old_text = change.get("old_text").and_then(serde_json::Value::as_str); + Some( + agent_client_protocol::schema::v1::Diff::new( + absolute_tool_path(path, cwd), + new_text.to_string(), + ) + .old_text(old_text.map(str::to_string)) + .into(), + ) +} + +fn tool_result_image(image: &serde_json::Value) -> Option { + let data = image.get("data_url")?.as_str()?; + let media_type = image.get("media_type")?.as_str()?; + let encoded = data + .strip_prefix("data:") + .and_then(|value| value.split_once(',')) + .map(|(_, encoded)| encoded) + .unwrap_or(data); + Some(ToolCallContent::from(ContentBlock::Image( + agent_client_protocol::schema::v1::ImageContent::new(encoded, media_type), + ))) +} + +pub(super) fn send_tool_call( + connection: &ConnectionTo, + session_id: &str, + tool_call: crate::llm::ToolCall, + cwd: &Path, +) -> Result<(), Error> { + let raw_input = serde_json::from_str(&tool_call.function.arguments) + .unwrap_or_else(|_| serde_json::json!({ "arguments": tool_call.function.arguments })); + let title = tool_title(&tool_call.function.name, &raw_input); + let update = SessionUpdate::ToolCall( + ToolCall::new(tool_call.id, title) + .kind(tool_kind(&tool_call.function.name)) + .status(ToolCallStatus::Pending) + .locations(tool_locations(&tool_call.function.name, &raw_input, cwd)) + .raw_input(raw_input), + ); + connection + .send_notification(SessionNotification::new(session_id.to_string(), update)) + .map_err(|_| internal_error()) +} + +pub(super) fn send_tool_result( + connection: &ConnectionTo, + session_id: &str, + result: crate::llm::ToolCallResult, + cwd: &Path, +) -> Result<(), Error> { + let payload = serde_json::from_str::(&result.content).unwrap_or_else( + |_| serde_json::json!({ "status": "error", "output_preview": result.content }), + ); + let status = match payload.get("status").and_then(serde_json::Value::as_str) { + Some("ok") => ToolCallStatus::Completed, + _ => ToolCallStatus::Failed, + }; + let content = tool_result_content(&payload, cwd); + let locations = tool_result_locations(&payload, cwd); + let mut fields = ToolCallUpdateFields::new() + .status(status) + .content((!content.is_empty()).then_some(content)) + .raw_output(payload); + if !locations.is_empty() { + fields = fields.locations(locations); + } + let update = SessionUpdate::ToolCallUpdate(ToolCallUpdate::new(result.tool_call_id, fields)); + connection + .send_notification(SessionNotification::new(session_id.to_string(), update)) + .map_err(|_| internal_error()) +} + +pub(super) fn replay_messages( + connection: &ConnectionTo, + session_id: &str, + messages: &[crate::session::types::Message], + cwd: &Path, +) -> Result<(), Error> { + for message in messages { + if crate::session::compaction::is_compaction_display_item(message) { + continue; + } + let message_id = message.id.clone(); + match message.role { + crate::session::types::MessageRole::User => { + if !message.content.is_empty() { + send_replay_text( + connection, + session_id, + &message_id, + &message.content, + true, + false, + )?; + } + } + crate::session::types::MessageRole::Assistant => { + for part in &message.parts { + match part.part_type.as_str() { + "text" => { + if let Some(text) = part.text_value() { + send_replay_text( + connection, + session_id, + &message_id, + text, + false, + false, + )?; + } + } + "reasoning" => { + if let Some(text) = part.text_value() { + send_replay_text( + connection, + session_id, + &message_id, + text, + false, + true, + )?; + } + } + "tool_call" => replay_tool_call(connection, session_id, part, cwd)?, + "tool_result" => replay_tool_result(connection, session_id, part, cwd)?, + _ => {} + } + } + } + crate::session::types::MessageRole::System + | crate::session::types::MessageRole::Tool => {} + } + } + Ok(()) +} + +pub(super) fn send_replay_text( + connection: &ConnectionTo, + session_id: &str, + message_id: &str, + text: &str, + user: bool, + thought: bool, +) -> Result<(), Error> { + let chunk = ContentChunk::new(text.to_string().into()).message_id(message_id); + let update = if user { + SessionUpdate::UserMessageChunk(chunk) + } else if thought { + SessionUpdate::AgentThoughtChunk(chunk) + } else { + SessionUpdate::AgentMessageChunk(chunk) + }; + connection + .send_notification(SessionNotification::new(session_id.to_string(), update)) + .map_err(|_| internal_error()) +} + +fn replay_tool_call( + connection: &ConnectionTo, + session_id: &str, + part: &crate::session::types::MessagePart, + cwd: &Path, +) -> Result<(), Error> { + let Some(tool_call_id) = part.tool_id() else { + return Ok(()); + }; + let name = part.tool_name().unwrap_or("tool"); + let input = part + .data + .get("args") + .cloned() + .unwrap_or(serde_json::Value::Null); + let status = match part.tool_status() { + Some("completed") | Some("ok") => ToolCallStatus::Completed, + Some("error") | Some("failed") => ToolCallStatus::Failed, + Some("running") => ToolCallStatus::InProgress, + _ => ToolCallStatus::Pending, + }; + let update = SessionUpdate::ToolCall( + ToolCall::new(tool_call_id.to_string(), tool_title(name, &input)) + .kind(tool_kind(name)) + .status(status) + .locations(tool_locations(name, &input, cwd)) + .raw_input(input), + ); + connection + .send_notification(SessionNotification::new(session_id.to_string(), update)) + .map_err(|_| internal_error()) +} + +fn replay_tool_result( + connection: &ConnectionTo, + session_id: &str, + part: &crate::session::types::MessagePart, + cwd: &Path, +) -> Result<(), Error> { + let Some(tool_call_id) = part.tool_id() else { + return Ok(()); + }; + let content = part + .data + .get("content") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string(); + send_tool_result( + connection, + session_id, + crate::llm::ToolCallResult { + tool_call_id: tool_call_id.to_string(), + role: "tool".to_string(), + name: part.tool_name().unwrap_or("tool").to_string(), + content, + }, + cwd, + ) +} + +pub(super) fn tool_kind(tool_name: &str) -> ToolKind { + match tool_name { + "bash" | "bash_output" | "bash_kill" | "bash_restart" | "terminal_session" => { + ToolKind::Execute + } + "webfetch" => ToolKind::Fetch, + "grep" | "glob" | "context" => ToolKind::Search, + "read" | "view_image" => ToolKind::Read, + "edit" | "write" | "write_files" | "apply_patch" => ToolKind::Edit, + "task" => ToolKind::Think, + _ => ToolKind::Other, + } +} + +fn tool_title(tool_name: &str, input: &serde_json::Value) -> String { + match tool_name { + "bash" => input + .get("command") + .and_then(serde_json::Value::as_str) + .unwrap_or(tool_name) + .to_string(), + "read" | "edit" | "write" | "grep" | "glob" => input + .get("filePath") + .or_else(|| input.get("filepath")) + .or_else(|| input.get("path")) + .or_else(|| input.get("pattern")) + .and_then(serde_json::Value::as_str) + .unwrap_or(tool_name) + .to_string(), + _ => tool_name.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use agent_client_protocol::schema::v1::ContentBlock; + use agent_client_protocol::schema::v1::ToolCallContent; + use agent_client_protocol::schema::v1::ToolKind; + use std::path::Path; + use std::path::PathBuf; + + #[test] + fn maps_crabcode_tools_to_acp_kinds() { + assert_eq!(tool_kind("bash"), ToolKind::Execute); + assert_eq!(tool_kind("read"), ToolKind::Read); + assert_eq!(tool_kind("apply_patch"), ToolKind::Edit); + assert_eq!(tool_kind("unknown"), ToolKind::Other); + } + + #[test] + fn builds_titles_from_tool_input() { + assert_eq!( + tool_title("bash", &serde_json::json!({"command": "cargo test"})), + "cargo test" + ); + assert_eq!( + tool_title("read", &serde_json::json!({"filePath": "src/main.rs"})), + "src/main.rs" + ); + } + + #[test] + fn acp_tool_result_prefers_full_output() { + let payload = serde_json::json!({ + "output": "complete tool output", + "output_preview": "short preview", + }); + + assert_eq!(tool_result_text(&payload), "complete tool output"); + } + + #[test] + fn acp_tool_result_supports_legacy_preview_payloads() { + let payload = serde_json::json!({"output_preview": "legacy output"}); + + assert_eq!(tool_result_text(&payload), "legacy output"); + } + + #[test] + fn acp_tool_locations_normalize_multi_file_and_patch_paths() { + let cwd = Path::new("/tmp/workspace"); + let write_locations = tool_locations( + "write_files", + &serde_json::json!({ + "files": [ + {"file_path": "src/a.rs", "content": "a"}, + {"file_path": "/tmp/b.rs", "content": "b"} + ] + }), + cwd, + ); + assert_eq!( + write_locations[0].path, + PathBuf::from("/tmp/workspace/src/a.rs") + ); + assert_eq!(write_locations[1].path, PathBuf::from("/tmp/b.rs")); + + let patch_locations = tool_locations( + "apply_patch", + &serde_json::json!({ + "patch": "*** Begin Patch\n*** Update File: src/a.rs\n*** Add File: src/b.rs\n*** End Patch" + }), + cwd, + ); + assert_eq!(patch_locations.len(), 2); + assert_eq!( + patch_locations[1].path, + PathBuf::from("/tmp/workspace/src/b.rs") + ); + } + + #[test] + fn acp_tool_result_emits_diff_location_and_image_content() { + let payload = serde_json::json!({ + "output": "updated", + "metadata": { + "path": "src/main.rs", + "line_number": 4, + "old_text": "fn old() {}", + "new_text": "fn new() {}" + }, + "images": [{ + "data_url": "data:image/png;base64,aGk=", + "media_type": "image/png" + }] + }); + let cwd = Path::new("/tmp/workspace"); + + let locations = tool_result_locations(&payload, cwd); + assert_eq!( + locations[0].path, + PathBuf::from("/tmp/workspace/src/main.rs") + ); + assert_eq!(locations[0].line, Some(4)); + + let content = tool_result_content(&payload, cwd); + let diff = content.iter().find_map(|item| match item { + ToolCallContent::Diff(diff) => Some(diff), + _ => None, + }); + let diff = diff.expect("diff content"); + assert_eq!(diff.path, PathBuf::from("/tmp/workspace/src/main.rs")); + assert_eq!(diff.old_text.as_deref(), Some("fn old() {}")); + assert_eq!(diff.new_text, "fn new() {}"); + assert!(content.iter().any(|item| matches!( + item, + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Image(image) if image.data == "aGk=" && image.mime_type == "image/png") + ))); + } + + #[test] + fn acp_tool_result_restores_native_mcp_resources() { + let metadata = serde_json::json!({ + "mcp_result": { + "content": [ + { + "type": "image", + "data": "aGk=", + "mimeType": "image/png", + "annotations": { "priority": 0.5 } + }, + { + "type": "resource_link", + "uri": "file:///tmp/readme.md", + "name": "readme", + "mimeType": "text/markdown", + "annotations": { "priority": 0.8 } + } + ] + } + }); + + let content = mcp_native_tool_content(&metadata); + assert_eq!(content.len(), 2); + assert!(matches!( + &content[0], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Image(image) + if image.data == "aGk=" && image.annotations.is_some()) + )); + assert!(matches!( + &content[1], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::ResourceLink(link) + if link.uri == "file:///tmp/readme.md") + )); + } + + #[test] + fn acp_tool_result_restores_mcp_audio_and_preserves_unknown_blocks() { + let metadata = serde_json::json!({ + "mcp_result": { + "content": [ + { "type": "audio", "data": "YXVkaW8=", "mimeType": "audio/wav" }, + { "type": "future_media", "payload": { "value": 1 } } + ] + } + }); + let content = mcp_native_tool_content(&metadata); + assert!(matches!( + &content[0], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Audio(audio) + if audio.data == "YXVkaW8=" && audio.mime_type == "audio/wav") + )); + assert!(matches!( + &content[1], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Text(text) + if text.text.contains("future_media") && text.text.contains("payload")) + )); + } +} diff --git a/src/acp/service/permissions.rs b/src/acp/service/permissions.rs new file mode 100644 index 0000000..ee620ec --- /dev/null +++ b/src/acp/service/permissions.rs @@ -0,0 +1,404 @@ +use super::notifications::absolute_tool_path; +use super::notifications::tool_kind; +use super::notifications::tool_locations; +use super::AcpSession; +use crate::config::configuration::LoadedConfig; +use agent_client_protocol::schema::v1::PermissionOption; +use agent_client_protocol::schema::v1::PermissionOptionKind; +use agent_client_protocol::schema::v1::RequestPermissionOutcome; +use agent_client_protocol::schema::v1::RequestPermissionRequest; +use agent_client_protocol::schema::v1::ToolCallContent; +use agent_client_protocol::schema::v1::ToolCallStatus; +use agent_client_protocol::schema::v1::ToolCallUpdate; +use agent_client_protocol::schema::v1::ToolCallUpdateFields; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use std::path::Path; +use std::path::PathBuf; + +fn permission_options() -> Vec { + vec![ + PermissionOption::new("once", "Allow once", PermissionOptionKind::AllowOnce), + PermissionOption::new( + "always", + "Always allow for this session", + PermissionOptionKind::AllowAlways, + ), + PermissionOption::new("reject", "Reject", PermissionOptionKind::RejectOnce), + ] +} + +fn permission_tool_content( + prompt: &crate::tools::PermissionPrompt, + cwd: &Path, +) -> Vec { + match prompt.tool_id.as_str() { + "write" => prompt + .raw_input + .get("file_path") + .or_else(|| prompt.raw_input.get("filePath")) + .and_then(serde_json::Value::as_str) + .zip( + prompt + .raw_input + .get("content") + .and_then(serde_json::Value::as_str), + ) + .map(|(path, new_text)| vec![preflight_diff(path, new_text, cwd)]) + .unwrap_or_default(), + "write_files" => prompt + .raw_input + .get("files") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|file| { + let path = file.get("file_path")?.as_str()?; + let content = file.get("content")?.as_str()?; + Some(preflight_diff(path, content, cwd)) + }) + .collect(), + "edit" => permission_edit_diff(&prompt.raw_input, cwd) + .into_iter() + .collect(), + "apply_patch" + if crate::tools::patch::patch_paths_from_params(&prompt.raw_input) + .iter() + .all(|path| preview_read_allowed(&absolute_tool_path(path, cwd), cwd)) => + { + crate::tools::patch::preview_patch(&prompt.raw_input, cwd) + .unwrap_or_default() + .into_iter() + .map(|change| { + agent_client_protocol::schema::v1::Diff::new(change.path, change.new_text) + .old_text(change.old_text) + .into() + }) + .collect() + } + _ => Vec::new(), + } +} + +// A write/edit approval is not a read approval. Keep ordinary workspace previews, +// but never acquire protected file contents just to display a permission request. +fn preview_read_allowed(path: &Path, cwd: &Path) -> bool { + !crate::tools::permission::is_outside_workdir(path, cwd) + && !crate::tools::permission::is_sensitive_path(path) + && !std::fs::canonicalize(path) + .ok() + .is_some_and(|resolved| crate::tools::permission::is_sensitive_path(&resolved)) +} + +fn preflight_diff(path: &str, new_text: &str, cwd: &Path) -> ToolCallContent { + let path = absolute_tool_path(path, cwd); + let old_text = preview_read_allowed(&path, cwd) + .then(|| std::fs::read_to_string(&path).ok()) + .flatten(); + agent_client_protocol::schema::v1::Diff::new(path, new_text.to_string()) + .old_text(old_text) + .into() +} + +fn permission_edit_diff(input: &serde_json::Value, cwd: &Path) -> Option { + let path = input + .get("file_path") + .or_else(|| input.get("filePath"))? + .as_str()?; + let old_string = input.get("old_string")?.as_str()?; + let new_string = input.get("new_string")?.as_str()?; + let absolute = absolute_tool_path(path, cwd); + if !preview_read_allowed(&absolute, cwd) { + return None; + } + let old_text = std::fs::read_to_string(&absolute).ok()?; + let new_text = if input + .get("replace_all") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) + { + old_text.replace(old_string, new_string) + } else { + old_text.replacen(old_string, new_string, 1) + }; + Some( + agent_client_protocol::schema::v1::Diff::new(absolute, new_text) + .old_text(old_text) + .into(), + ) +} + +fn permission_tool_call_id(tool_call_id: Option<&str>) -> String { + tool_call_id + .map(str::to_string) + .unwrap_or_else(|| format!("permission:{}", cuid2::create_id())) +} + +pub(super) fn configured_tool_permissions( + cwd: &Path, + config: &LoadedConfig, +) -> crate::tools::ToolPermissions { + let mut policies = crate::tools::AgentToolPolicies::default(); + for (mode, tools) in config.merged_config.agent_registry.tool_policy_map() { + policies = policies.with_custom_tools(mode, tools); + } + crate::tools::ToolPermissions::new(cwd) + .with_agent_policies(policies) + .with_permission_rules(config.merged_config.permission_rules.clone()) + .with_agent_permission_rules(config.merged_config.agent_registry.permission_rules_map()) +} + +pub(super) fn tool_permissions(session: &AcpSession) -> crate::tools::ToolPermissions { + session.tool_permissions.clone() +} + +pub(super) async fn request_permission( + connection: &ConnectionTo, + session_id: &str, + prompt: &crate::tools::PermissionPrompt, +) -> crate::tools::PermissionResponse { + let tool_call_id = permission_tool_call_id(prompt.tool_call_id.as_deref()); + let input = serde_json::json!({ + "tool": prompt.tool_id, + "permission": prompt.permission, + "patterns": prompt.patterns, + "target": prompt.target, + "command": prompt.command, + "workdir": prompt.workdir, + "reason": prompt.reason, + "input": prompt.raw_input, + }); + let cwd = PathBuf::from(&prompt.workspace); + let content = permission_tool_content(prompt, &cwd); + let tool_call = ToolCallUpdate::new(tool_call_id, { + let mut fields = ToolCallUpdateFields::new() + .title(permission_title(prompt)) + .kind(tool_kind(&prompt.tool_id)) + .status(ToolCallStatus::Pending) + .locations(tool_locations(&prompt.tool_id, &prompt.raw_input, &cwd)) + .raw_input(input); + if !content.is_empty() { + fields = fields.content(content); + } + fields + }); + let request = + RequestPermissionRequest::new(session_id.to_string(), tool_call, permission_options()); + let Ok(response) = connection.send_request(request).block_task().await else { + return crate::tools::PermissionResponse::Deny; + }; + match response.outcome { + RequestPermissionOutcome::Selected(selection) + if selection.option_id.to_string() == "once" => + { + crate::tools::PermissionResponse::AllowOnce + } + RequestPermissionOutcome::Selected(selection) + if selection.option_id.to_string() == "always" => + { + crate::tools::PermissionResponse::AllowAlways + } + RequestPermissionOutcome::Selected(_) | RequestPermissionOutcome::Cancelled => { + crate::tools::PermissionResponse::Deny + } + _ => crate::tools::PermissionResponse::Deny, + } +} + +fn permission_title(prompt: &crate::tools::PermissionPrompt) -> String { + prompt + .command + .as_deref() + .or(prompt.target.as_deref()) + .unwrap_or(&prompt.reason) + .to_string() +} + +#[cfg(test)] +mod tests { + use super::super::test_support::test_session; + use super::*; + use agent_client_protocol::schema::v1::PermissionOptionKind; + use agent_client_protocol::schema::v1::ToolCallContent; + + #[test] + fn acp_permission_prefers_originating_tool_call_id() { + assert_eq!(permission_tool_call_id(Some("call_123")), "call_123"); + } + + #[test] + fn acp_permission_generates_fallback_id_without_origin() { + assert!(permission_tool_call_id(None).starts_with("permission:")); + } + + #[test] + fn acp_permission_offers_explicit_session_wide_allow() { + let options = permission_options(); + let always = options + .iter() + .find(|option| option.option_id.to_string() == "always") + .expect("always option"); + assert_eq!(always.name, "Always allow for this session"); + assert_eq!(always.kind, PermissionOptionKind::AllowAlways); + } + + #[tokio::test] + async fn acp_session_reuses_permission_grants_across_turns() { + let temp = tempfile::tempdir().expect("temp dir"); + let workspace = temp.path().join("workspace"); + let external = temp.path().join("external"); + std::fs::create_dir_all(&workspace).expect("workspace"); + std::fs::create_dir_all(&external).expect("external"); + + let mut session = test_session(); + session.cwd = workspace.clone(); + session.tool_permissions = crate::tools::ToolPermissions::new(&workspace); + let first_turn = tool_permissions(&session); + let second_turn = tool_permissions(&session); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let params = serde_json::json!({ "file_path": external.join("one.txt") }); + + let pending = tokio::spawn({ + let permissions = first_turn.clone(); + let params = params.clone(); + let tx = tx.clone(); + async move { + permissions + .preflight("build", "read", ¶ms, Some(&tx)) + .await + } + }); + let prompt = match rx.recv().await { + Some(crate::llm::ChunkMessage::PermissionRequest(prompt)) => prompt, + _ => panic!("expected permission request"), + }; + let _ = prompt + .response_tx + .send(crate::tools::PermissionResponse::AllowAlways); + assert!(pending.await.expect("permission task").is_ok()); + + let next_params = serde_json::json!({ "file_path": external.join("nested/two.txt") }); + assert!(second_turn + .preflight("build", "read", &next_params, Some(&tx)) + .await + .is_ok()); + assert!(rx.try_recv().is_err()); + } + + #[test] + fn acp_permission_edit_includes_preflight_diff() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("main.rs"); + std::fs::write(&path, "fn old() {}\n").unwrap(); + let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); + let prompt = crate::tools::PermissionPrompt { + tool_call_id: Some("call_edit".to_string()), + tool_id: "edit".to_string(), + action: crate::tools::PermissionAction::Write, + permission: "edit".to_string(), + patterns: vec![path.to_string_lossy().into_owned()], + target: Some(path.to_string_lossy().into_owned()), + command: None, + workdir: None, + workspace: dir.path().to_string_lossy().into_owned(), + reason: "approval".to_string(), + raw_input: serde_json::json!({ + "file_path": path, + "old_string": "old", + "new_string": "new", + "replace_all": false + }), + response_tx, + }; + + let content = permission_tool_content(&prompt, dir.path()); + let ToolCallContent::Diff(diff) = &content[0] else { + panic!("expected preflight diff"); + }; + assert_eq!(diff.old_text.as_deref(), Some("fn old() {}\n")); + assert_eq!(diff.new_text, "fn new() {}\n"); + } + + #[test] + fn acp_permission_apply_patch_includes_preflight_diff() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("note.txt"); + std::fs::write(&path, "before\n").unwrap(); + let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); + let prompt = crate::tools::PermissionPrompt { + tool_call_id: Some("call_patch".to_string()), + tool_id: "apply_patch".to_string(), + action: crate::tools::PermissionAction::Write, + permission: "edit".to_string(), + patterns: vec![path.to_string_lossy().into_owned()], + target: Some(path.to_string_lossy().into_owned()), + command: None, + workdir: None, + workspace: dir.path().to_string_lossy().into_owned(), + reason: "approval".to_string(), + raw_input: serde_json::json!({ + "patch": format!( + "*** Begin Patch\n*** Update File: {}\n@@\n-before\n+after\n*** End Patch\n", + path.display() + ) + }), + response_tx, + }; + + let content = permission_tool_content(&prompt, dir.path()); + let ToolCallContent::Diff(diff) = &content[0] else { + panic!("expected apply_patch preflight diff"); + }; + assert_eq!(diff.path, path); + assert_eq!(diff.old_text.as_deref(), Some("before\n")); + assert_eq!(diff.new_text, "after\n"); + } + + #[test] + fn previews_keep_workspace_diffs_but_do_not_read_protected_files() { + let workspace = tempfile::tempdir().unwrap(); + let external = tempfile::tempdir().unwrap(); + let ordinary = workspace.path().join("hello.txt"); + std::fs::write(&ordinary, "before").unwrap(); + let preview = + serde_json::to_string(&preflight_diff("hello.txt", "after", workspace.path())).unwrap(); + assert!(preview.contains("before")); + for path in [ + workspace.path().join(".env"), + external.path().join("secret.txt"), + ] { + std::fs::write(&path, "secret-never-preview").unwrap(); + let preview = serde_json::to_string(&preflight_diff( + path.to_str().unwrap(), + "after", + workspace.path(), + )) + .unwrap(); + assert!(!preview.contains("secret-never-preview")); + assert!(permission_edit_diff( + &serde_json::json!({ + "file_path": path, "old_string": "secret", "new_string": "public" + }), + workspace.path() + ) + .is_none()); + } + } + + #[cfg(unix)] + #[test] + fn previews_reject_symlink_escape_and_sensitive_alias() { + let workspace = tempfile::tempdir().unwrap(); + let external = tempfile::tempdir().unwrap(); + let secret = external.path().join("secret.txt"); + std::fs::write(&secret, "secret").unwrap(); + let link = workspace.path().join("alias.txt"); + std::os::unix::fs::symlink(secret, &link).unwrap(); + assert!(!preview_read_allowed(&link, workspace.path())); + let env = workspace.path().join(".env"); + std::fs::write(&env, "secret").unwrap(); + let alias = workspace.path().join("config.txt"); + std::os::unix::fs::symlink(env, &alias).unwrap(); + assert!(!preview_read_allowed(&alias, workspace.path())); + } +} diff --git a/src/acp/service/prompt.rs b/src/acp/service/prompt.rs new file mode 100644 index 0000000..b8763dc --- /dev/null +++ b/src/acp/service/prompt.rs @@ -0,0 +1,367 @@ +use super::attachments::model_supports_audio; +use super::attachments::prompt_content; +use super::attachments::ManagedAttachmentGuard; +use super::cancellable; +use super::commands::command_text; +use super::commands::compact_command; +use super::commands::expand_slash_command; +use super::commands::SlashExpansion; +use super::config::find_selectable_model; +use super::config::model_context_window; +use super::config::resolved_reasoning; +use super::internal_error; +use super::notifications::send_replay_text; +use super::permissions::tool_permissions; +use super::usage::acp_context_tokens; +use super::usage::acp_session_cost; +use super::usage::send_usage; +use super::AcpService; +use agent_client_protocol::schema::v1::ContentBlock; +use agent_client_protocol::schema::v1::PromptResponse; +use agent_client_protocol::schema::v1::StopReason; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use agent_client_protocol::Error; +use tokio_util::sync::CancellationToken; + +pub(super) fn acp_request_messages( + history: &[crate::session::types::Message], + system_message: crate::session::types::Message, +) -> Vec { + let mut messages = crate::session::compaction::filter_messages_for_context(history); + messages.insert(0, system_message); + messages +} + +impl AcpService { + pub async fn prompt( + &self, + session_id: String, + prompt: Vec, + connection: ConnectionTo, + ) -> Result { + let cancellation = self.sessions.lock().await.begin_prompt(&session_id)?; + // Only individual blocking operations are cancellable. Dropping prompt_inner + // would skip transcript persistence and leave the session marked active. + let mut result = self + .prompt_inner(session_id.clone(), prompt, connection, cancellation.clone()) + .await; + if cancellation.is_cancelled() { + result = Ok(PromptResponse::new(StopReason::Cancelled)); + } + // Cover setup errors and local commands as well as stream/compaction exits. + if let Ok(mut manager) = self.session_manager.lock() { + let status = if cancellation.is_cancelled() + || result + .as_ref() + .is_ok_and(|response| response.stop_reason == StopReason::Cancelled) + { + crate::session::types::SessionStatus::Interrupted + } else if result.is_err() { + crate::session::types::SessionStatus::Failed + } else { + crate::session::types::SessionStatus::Idle + }; + if manager.get_session_ref(&session_id).is_some() { + if manager + .set_session_status(&session_id, status, None) + .is_err() + { + result = Err(internal_error()); + } + } + } else { + result = Err(internal_error()); + } + // Also stop any detached producer after an error or early return. + cancellation.cancel(); + self.sessions.lock().await.finish_prompt(&session_id); + result + } + + async fn prompt_inner( + &self, + session_id: String, + prompt: Vec, + connection: ConnectionTo, + cancellation: CancellationToken, + ) -> Result { + let mut session = self + .sessions + .lock() + .await + .attachments + .get(&session_id) + .cloned() + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + let compact_text = command_text(&prompt); + if compact_command(&compact_text)? { + if prompt + .iter() + .any(|part| !matches!(part, ContentBlock::Text(_))) + { + return Err(Error::invalid_params().data("/compact does not accept attachments")); + } + return self + .compact_session(&session_id, session, connection, cancellation) + .await; + } + let expansion = if compact_text.trim_start().starts_with('/') { + let Some(expansion) = + cancellable(&cancellation, expand_slash_command(&session, &compact_text)).await + else { + return Ok(PromptResponse::new(StopReason::Cancelled)); + }; + Some(expansion?) + } else { + None + }; + let expanded_prompt = match expansion { + Some(SlashExpansion::LocalResult(text)) => { + if prompt + .iter() + .any(|part| !matches!(part, ContentBlock::Text(_))) + { + return Err(Error::invalid_params() + .data("local ACP commands do not accept attachments")); + } + let message_id = cuid2::create_id(); + send_replay_text(&connection, &session_id, &message_id, &text, false, false)?; + return Ok(PromptResponse::new(StopReason::EndTurn)); + } + Some(SlashExpansion::Prompt { + prompt, + agent, + model, + }) => { + if let Some(agent) = agent { + if session + .config + .merged_config + .agent_registry + .get(&agent) + .is_none() + { + return Err(Error::invalid_params() + .data(format!("custom command references unknown agent: {agent}"))); + } + session.agent = agent; + } + if let Some(model_ref) = model { + let (provider, model) = crate::app::parse_model_ref(&model_ref); + let canonical = format!("{provider}/{model}"); + let model = find_selectable_model(&session.models, &canonical)?; + session.provider.clone_from(&model.provider_id); + session.model.clone_from(&model.id); + session.reasoning = resolved_reasoning(&session, session.reasoning_selection); + session.context_window = model_context_window( + &session.config, + &session.models, + &session.provider, + &session.model, + ); + } + Some(prompt) + } + Some(SlashExpansion::Btw(question)) => { + if prompt + .iter() + .any(|part| !matches!(part, ContentBlock::Text(_))) + { + return Err(Error::invalid_params().data("/btw does not accept attachments")); + } + return self + .btw_session(&session_id, session, question, connection, cancellation) + .await; + } + None => None, + }; + let supports_images = session + .models + .iter() + .find(|model| model.provider_id == session.provider && model.id == session.model) + .is_some_and(|model| model.attachment); + let supports_audio = + model_supports_audio(&session.config, &session.provider, &session.model); + let (prompt, local_image_paths, local_audio_paths) = prompt_content( + prompt, + expanded_prompt.as_deref(), + supports_images, + supports_audio, + &session_id, + &session, + )?; + let mut managed_paths = local_image_paths.clone(); + managed_paths.extend(local_audio_paths.clone()); + let mut attachment_guard = ManagedAttachmentGuard::new(managed_paths); + if prompt.trim().is_empty() { + return Err(Error::invalid_params().data("prompt must include text content")); + } + let mut messages = { + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + let stored = manager + .get_session(&session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + stored.messages.clone() + }; + let mut user_message = crate::session::types::Message::user(&prompt); + user_message.local_image_paths = local_image_paths; + user_message.local_audio_paths = local_audio_paths; + user_message.provider = Some(session.provider.clone()); + user_message.model = Some(session.model.clone()); + user_message.agent_mode = Some(session.agent.clone()); + let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); + let process_registry = std::sync::Arc::new(crate::tools::ProcessRegistry::new()); + let prompt_registry = cancellable( + &cancellation, + crate::tools::initialize_tool_registry_with_dynamic_config( + Some(sender.clone()), + tool_permissions(&session), + session.config.merged_config.agent_registry.clone(), + cancellation.clone(), + Some(&session.provider), + &session.config.merged_config.websearch, + &session.config.merged_config.mcp, + &session.cwd, + process_registry.clone(), + ), + ) + .await; + let Some(prompt_registry) = prompt_registry else { + return Ok(PromptResponse::new(StopReason::Cancelled)); + }; + let is_git_repo = + crate::utils::git::is_git_repo(&session.cwd.to_string_lossy()).unwrap_or(false); + let system_prompt = cancellable( + &cancellation, + crate::prompt::SystemPromptComposer::new( + &session.model, + session.cwd.to_string_lossy(), + is_git_repo, + std::env::consts::OS, + ) + .with_tool_registry(prompt_registry.clone()) + .with_agent_registry(session.config.merged_config.agent_registry.clone()) + .with_active_agent(session.agent.clone()) + .compose(), + ) + .await; + let Some(system_prompt) = system_prompt else { + return Ok(PromptResponse::new(StopReason::Cancelled)); + }; + let system_message = crate::session::types::Message::system(system_prompt); + let prompt_tokens = crate::session::compaction::message_context_tokens(&system_message); + let auto_compacted = self + .maybe_auto_compact_session( + &session_id, + &session, + &connection, + cancellation.clone(), + prompt_tokens, + crate::session::compaction::message_context_tokens(&user_message), + ) + .await?; + if cancellation.is_cancelled() { + return Ok(PromptResponse::new(StopReason::Cancelled)); + } + if auto_compacted { + messages = self + .session_manager + .lock() + .map_err(|_| internal_error())? + .get_session_ref(&session_id) + .map(|stored| stored.messages.clone()) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + } + { + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + manager + .add_message_to_session(&session_id, &user_message) + .map_err(|_| internal_error())?; + attachment_guard.commit(); + manager + .set_session_status( + &session_id, + crate::session::types::SessionStatus::Streaming, + None, + ) + .map_err(|_| internal_error())?; + } + messages.push(user_message); + let base_context_tokens = acp_context_tokens(&messages, prompt_tokens); + let base_cost = acp_session_cost(&messages); + messages = acp_request_messages(&messages, system_message); + send_usage( + &connection, + &session_id, + &session, + base_context_tokens, + (base_cost > 0.0).then_some(base_cost), + None, + )?; + + let stream_session_id = session_id.clone(); + let stream_cancellation = cancellation.clone(); + let stream_sender = sender.clone(); + let stream_session = session.clone(); + let stream_tool_registry = prompt_registry; + let stream_process_registry = process_registry; + tokio::spawn(async move { + let result = crate::llm::client::stream_llm_with_cancellation( + stream_cancellation, + stream_session_id, + stream_session.provider.clone(), + stream_session.model.clone(), + stream_session.reasoning, + stream_session.agent.clone(), + stream_session + .config + .merged_config + .agent_registry + .get(&stream_session.agent) + .and_then(|agent| agent.max_steps), + stream_session.config.merged_config.agent_registry.clone(), + tool_permissions(&stream_session), + stream_session.config.merged_config.websearch.clone(), + stream_session.config.merged_config.mcp.clone(), + stream_session.config.merged_config.compaction.clone(), + stream_session.cwd.to_string_lossy().to_string(), + Some(stream_tool_registry), + messages, + sender, + stream_process_registry, + ) + .await; + if let Err(error) = result { + let _ = stream_sender.send(crate::llm::ChunkMessage::Failed(error.to_string())); + } + let _ = stream_sender.send(crate::llm::ChunkMessage::End); + }); + + let response = self + .receive_prompt_stream( + &session_id, + &session, + &connection, + &cancellation, + receiver, + base_context_tokens, + prompt_tokens, + base_cost, + ) + .await?; + if !cancellation.is_cancelled() && response.stop_reason != StopReason::Cancelled { + let _ = self + .maybe_auto_compact_session( + &session_id, + &session, + &connection, + cancellation.clone(), + prompt_tokens, + 0, + ) + .await?; + } + Ok(response) + } +} diff --git a/src/acp/service/questions.rs b/src/acp/service/questions.rs new file mode 100644 index 0000000..43d9b4d --- /dev/null +++ b/src/acp/service/questions.rs @@ -0,0 +1,337 @@ +use agent_client_protocol::schema::v1::CreateElicitationRequest; +use agent_client_protocol::schema::v1::ElicitationAction; +use agent_client_protocol::schema::v1::ElicitationContentValue; +use agent_client_protocol::schema::v1::ElicitationFormMode; +use agent_client_protocol::schema::v1::ElicitationSchema; +use agent_client_protocol::schema::v1::ElicitationSessionScope; +use agent_client_protocol::schema::v1::EnumOption; +use agent_client_protocol::schema::v1::MultiSelectPropertySchema; +use agent_client_protocol::schema::v1::StringPropertySchema; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use std::collections::HashMap; +use tokio_util::sync::CancellationToken; + +struct AcpQuestionField { + selection: String, + custom: String, + labels: HashMap, + multiple: bool, +} + +struct AcpQuestionForm { + request: CreateElicitationRequest, + fields: Vec, +} + +pub(super) fn skipped_question_answers(questions: &serde_json::Value) -> serde_json::Value { + let count = questions.as_array().map_or(1, Vec::len); + serde_json::Value::Array( + (0..count) + .map(|_| serde_json::Value::Array(Vec::new())) + .collect(), + ) +} + +fn question_text(question: &serde_json::Value, key: &str, fallback: &str) -> String { + question + .get(key) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(fallback) + .to_string() +} + +fn acp_question_form( + session_id: &str, + tool_call_id: Option<&str>, + questions: &serde_json::Value, +) -> AcpQuestionForm { + let question_items = questions + .as_array() + .cloned() + .unwrap_or_else(|| vec![questions.clone()]); + let mut schema = ElicitationSchema::new() + .title("Agent questions") + .description("Answer any fields you want; blank fields are treated as skipped."); + let mut fields = Vec::with_capacity(question_items.len()); + + for (question_index, question) in question_items.iter().enumerate() { + let selection = format!("question_{question_index}"); + let custom = format!("question_{question_index}_custom"); + let prompt = question_text(question, "question", "Question"); + let header = question_text( + question, + "header", + &format!("Question {}", question_index + 1), + ); + let mut labels = HashMap::new(); + let options = question + .get("options") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .enumerate() + .filter_map(|(option_index, option)| { + let label = option + .get("label") + .and_then(serde_json::Value::as_str) + .or_else(|| option.as_str())? + .trim(); + if label.is_empty() { + return None; + } + + let value = format!("q{question_index}_option_{option_index}"); + labels.insert(value.clone(), label.to_string()); + let description = option + .get("description") + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()); + let mut option = EnumOption::new(value, label); + if let Some(description) = description { + option = option.description(description); + } + Some(option) + }) + .collect::>(); + let multiple = question + .get("multiple") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + + if multiple { + schema = schema.property( + selection.clone(), + MultiSelectPropertySchema::titled(options) + .title(header.clone()) + .description(prompt.clone()), + false, + ); + } else { + schema = schema.property( + selection.clone(), + StringPropertySchema::new() + .title(header.clone()) + .description(prompt.clone()) + .one_of(options), + false, + ); + } + schema = schema.property( + custom.clone(), + StringPropertySchema::new() + .title(format!("{header}: custom answer")) + .description("Optional free-form answer."), + false, + ); + fields.push(AcpQuestionField { + selection, + custom, + labels, + multiple, + }); + } + + let scope = ElicitationSessionScope::new(session_id.to_string()) + .tool_call_id(tool_call_id.map(agent_client_protocol::schema::v1::ToolCallId::new)); + let request = CreateElicitationRequest::new( + ElicitationFormMode::new(scope, schema), + "The agent needs additional input to continue.", + ); + AcpQuestionForm { request, fields } +} + +fn acp_question_answers( + fields: &[AcpQuestionField], + action: ElicitationAction, +) -> serde_json::Value { + let ElicitationAction::Accept(accepted) = action else { + return serde_json::Value::Array( + fields + .iter() + .map(|_| serde_json::Value::Array(Vec::new())) + .collect(), + ); + }; + let content = accepted.content.unwrap_or_default(); + serde_json::Value::Array( + fields + .iter() + .map(|field| { + let mut answers = Vec::new(); + match content.get(&field.selection) { + Some(ElicitationContentValue::String(value)) if !field.multiple => { + if let Some(label) = field.labels.get(value) { + answers.push(serde_json::Value::String(label.clone())); + } + } + Some(ElicitationContentValue::StringArray(values)) if field.multiple => { + for value in values { + if let Some(label) = field.labels.get(value) { + if !answers.iter().any(|answer| answer.as_str() == Some(label)) { + answers.push(serde_json::Value::String(label.clone())); + } + } + } + } + _ => {} + } + if let Some(ElicitationContentValue::String(custom)) = content.get(&field.custom) { + let custom = custom.trim(); + if !custom.is_empty() { + if !field.multiple { + answers.clear(); + } + answers.push(serde_json::Value::String( + custom.chars().take(8_192).collect(), + )); + } + } + serde_json::Value::Array(answers) + }) + .collect(), + ) +} + +pub(super) async fn request_questions( + connection: &ConnectionTo, + session_id: &str, + tool_call_id: Option<&str>, + questions: &serde_json::Value, + cancellation: &CancellationToken, +) -> serde_json::Value { + let form = acp_question_form(session_id, tool_call_id, questions); + let request = connection.send_request(form.request).block_task(); + tokio::pin!(request); + let response = tokio::select! { + _ = cancellation.cancelled() => return skipped_question_answers(questions), + response = &mut request => response, + }; + let Ok(response) = response else { + return skipped_question_answers(questions); + }; + acp_question_answers(&form.fields, response.action) +} + +#[cfg(test)] +mod tests { + use super::*; + use agent_client_protocol::schema::v1::ElicitationAction; + use agent_client_protocol::schema::v1::ElicitationContentValue; + + #[test] + fn acp_question_form_preserves_single_multi_custom_and_scope() { + let form = acp_question_form( + "session_1", + Some("question_call_1"), + &serde_json::json!([ + { + "question": "Pick one", + "header": "Single", + "options": [ + {"label": "A", "description": "First"}, + {"label": "B", "description": "Second"} + ] + }, + { + "question": "Pick several", + "header": "Multiple", + "multiple": true, + "options": [ + {"label": "X", "description": "First"}, + {"label": "Y", "description": "Second"} + ] + } + ]), + ); + let wire = serde_json::to_value(&form.request).expect("elicitation request"); + + assert_eq!(wire["mode"], "form"); + assert_eq!(wire["sessionId"], "session_1"); + assert_eq!(wire["toolCallId"], "question_call_1"); + assert_eq!( + wire["requestedSchema"]["properties"]["question_0"]["type"], + "string" + ); + assert_eq!( + wire["requestedSchema"]["properties"]["question_0"]["oneOf"][0]["title"], + "A" + ); + assert_eq!( + wire["requestedSchema"]["properties"]["question_1"]["type"], + "array" + ); + assert_eq!( + wire["requestedSchema"]["properties"]["question_1_custom"]["type"], + "string" + ); + } + + #[test] + fn acp_question_answers_restore_labels_and_custom_text() { + let form = acp_question_form( + "session_1", + None, + &serde_json::json!([ + { + "question": "Pick one", + "options": [{"label": "A"}, {"label": "B"}] + }, + { + "question": "Pick several", + "multiple": true, + "options": [{"label": "X"}, {"label": "Y"}] + } + ]), + ); + let mut content = std::collections::BTreeMap::new(); + content.insert( + "question_0".to_string(), + ElicitationContentValue::String("q0_option_1".to_string()), + ); + content.insert( + "question_1".to_string(), + ElicitationContentValue::StringArray(vec![ + "q1_option_0".to_string(), + "q1_option_1".to_string(), + ]), + ); + content.insert( + "question_1_custom".to_string(), + ElicitationContentValue::String("Other choice".to_string()), + ); + let action = ElicitationAction::Accept( + agent_client_protocol::schema::v1::ElicitationAcceptAction::new().content(content), + ); + + assert_eq!( + acp_question_answers(&form.fields, action), + serde_json::json!([["B"], ["X", "Y", "Other choice"]]) + ); + assert_eq!( + acp_question_answers(&form.fields, ElicitationAction::Cancel), + serde_json::json!([[], []]) + ); + + let mut custom_content = std::collections::BTreeMap::new(); + custom_content.insert( + "question_0".to_string(), + ElicitationContentValue::String("q0_option_0".to_string()), + ); + custom_content.insert( + "question_0_custom".to_string(), + ElicitationContentValue::String("Custom only".to_string()), + ); + let custom_action = ElicitationAction::Accept( + agent_client_protocol::schema::v1::ElicitationAcceptAction::new() + .content(custom_content), + ); + assert_eq!( + acp_question_answers(&form.fields, custom_action), + serde_json::json!([["Custom only"], []]) + ); + } +} diff --git a/src/acp/service/sessions.rs b/src/acp/service/sessions.rs new file mode 100644 index 0000000..65542b3 --- /dev/null +++ b/src/acp/service/sessions.rs @@ -0,0 +1,411 @@ +use super::config::merge_acp_mcp_servers; +use super::config::model_context_window; +use super::config::model_reasoning; +use super::config::resolve_model; +use super::config::session_config_options; +use super::config::session_modes; +use super::internal_error; +use super::notifications::replay_messages; +use super::permissions::configured_tool_permissions; +use super::AcpService; +use super::AcpSession; +use agent_client_protocol::schema::v1::ListSessionsResponse; +use agent_client_protocol::schema::v1::LoadSessionResponse; +use agent_client_protocol::schema::v1::McpServer; +use agent_client_protocol::schema::v1::NewSessionResponse; +use agent_client_protocol::schema::v1::ResumeSessionResponse; +use agent_client_protocol::schema::v1::SessionInfo; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use agent_client_protocol::Error; +use std::path::Path; +use std::path::PathBuf; + +fn acp_session_info( + session: crate::session::manager::SessionInfo, + root_id: Option, +) -> SessionInfo { + let mut meta = serde_json::Map::new(); + meta.insert( + "crabcode".to_string(), + serde_json::json!({ + "parentSessionId": session.parent_id, + "rootSessionId": root_id, + }), + ); + SessionInfo::new(session.id, session.workspace_path) + .title(session.title) + .updated_at(system_time_to_iso8601(session.updated_at)) + .meta(meta) +} + +fn session_page( + cursor: Option<&str>, + total: usize, +) -> Result<(usize, usize, Option), Error> { + let offset = cursor + .unwrap_or("0") + .parse::() + .map_err(|_| Error::invalid_params().data("invalid session list cursor"))?; + let end = offset.saturating_add(100).min(total); + let next_cursor = (end < total).then(|| end.to_string()); + Ok((offset.min(total), end, next_cursor)) +} + +fn workspace_path(path: &Path) -> Result { + if !path.is_absolute() { + return Err(Error::invalid_params().data("cwd must be an absolute path")); + } + let path = path + .canonicalize() + .map_err(|_| Error::resource_not_found(Some(path.to_string_lossy().to_string())))?; + if !path.is_dir() { + return Err(Error::invalid_params().data("cwd must be a directory")); + } + Ok(path) +} + +fn system_time_to_iso8601(value: std::time::SystemTime) -> String { + chrono::DateTime::::from(value).to_rfc3339() +} + +impl AcpService { + pub async fn new_session( + &self, + cwd: PathBuf, + mcp_servers: Vec, + ) -> Result { + let cwd = workspace_path(&cwd)?; + let mut config = + crate::config::ConfigLoader::load_for(&cwd).map_err(|_| internal_error())?; + merge_acp_mcp_servers(&mut config, mcp_servers); + crate::skill::init_skill_store(&config.xdg_config_home, &config.project_root); + let (provider, model) = resolve_model(&config); + let models = crate::model::catalog::selectable_models(&config, None) + .await + .map_err(|_| internal_error())?; + let reasoning = model_reasoning(&config, &models, &provider, &model); + let reasoning_selection = + reasoning.unwrap_or(crate::model::reasoning::ReasoningEffort::None); + let context_window = model_context_window(&config, &models, &provider, &model); + let agent = config + .merged_config + .default_agent + .clone() + .unwrap_or_else(|| { + config + .merged_config + .agent_registry + .default_agent() + .to_string() + }); + + let session_id = { + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + manager + .switch_current_workspace_path(&cwd.to_string_lossy()) + .map_err(|_| internal_error())?; + manager.create_session(Some("New session".to_string())) + }; + + let skills = crate::skill::SkillStore::load(&config.xdg_config_home, &config.project_root); + let tool_permissions = configured_tool_permissions(&cwd, &config); + self.sessions.lock().await.attachments.insert( + session_id.clone(), + AcpSession { + cwd, + config, + tool_permissions, + skills, + models, + provider, + model, + agent, + reasoning_selection, + reasoning, + context_window, + cancellation: None, + }, + ); + + let session = self + .sessions + .lock() + .await + .attachments + .get(&session_id) + .cloned() + .ok_or_else(internal_error)?; + Ok(NewSessionResponse::new(session_id) + .modes(session_modes(&session)) + .config_options(session_config_options(&session))) + } + + pub async fn list_sessions( + &self, + cwd: Option, + cursor: Option, + ) -> Result { + let cwd = cwd.as_deref().map(workspace_path).transpose()?; + let manager = self.session_manager.lock().map_err(|_| internal_error())?; + let mut sessions = manager + .list_sessions() + .into_iter() + .filter(|session| session.archived_at.is_none()) + .filter(|session| { + cwd.as_ref() + .is_none_or(|cwd| session.workspace_path == cwd.to_string_lossy()) + }) + .collect::>(); + sessions.sort_by(|left, right| right.updated_at.cmp(&left.updated_at)); + + let (offset, end, next_cursor) = session_page(cursor.as_deref(), sessions.len())?; + Ok(ListSessionsResponse::new( + sessions + .into_iter() + .skip(offset) + .take(end.saturating_sub(offset)) + .map(|session| { + let root_id = manager.root_session_id_for(&session.id); + acp_session_info(session, root_id) + }) + .collect(), + ) + .next_cursor(next_cursor)) + } + + pub async fn load_session( + &self, + session_id: String, + cwd: PathBuf, + connection: ConnectionTo, + ) -> Result { + let (session, messages) = self.attach_persisted_session(&session_id, cwd).await?; + replay_messages(&connection, &session_id, &messages, &session.cwd)?; + Ok(LoadSessionResponse::new() + .modes(session_modes(&session)) + .config_options(session_config_options(&session))) + } + + pub async fn resume_session( + &self, + session_id: String, + cwd: PathBuf, + ) -> Result { + let (session, _) = self.attach_persisted_session(&session_id, cwd).await?; + Ok(ResumeSessionResponse::new() + .modes(session_modes(&session)) + .config_options(session_config_options(&session))) + } + + pub async fn fork_session( + &self, + session_id: String, + cwd: PathBuf, + ) -> Result { + let (source, messages) = self.attach_persisted_session(&session_id, cwd).await?; + let fork_id = { + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + let source_title = manager + .get_session(&session_id) + .map(|session| session.title.clone()) + .unwrap_or_else(|| session_id.clone()); + manager + .switch_current_workspace_path(&source.cwd.to_string_lossy()) + .map_err(|_| internal_error())?; + let fork_id = manager.create_session(Some(format!("{source_title} (fork)"))); + let messages = + match crate::persistence::attachments::clone_messages(&messages, &fork_id) { + Ok(messages) => messages, + Err(_) => { + manager.delete_session(&fork_id); + return Err(internal_error()); + } + }; + if manager + .replace_session_messages(&fork_id, messages) + .is_err() + { + manager.delete_session(&fork_id); + return Err(internal_error()); + } + fork_id + }; + self.sessions.lock().await.attachments.insert( + fork_id.clone(), + AcpSession { + tool_permissions: configured_tool_permissions(&source.cwd, &source.config), + ..source.clone() + }, + ); + Ok(NewSessionResponse::new(fork_id) + .modes(session_modes(&source)) + .config_options(session_config_options(&source))) + } + + pub async fn close_session(&self, session_id: &str) { + if let Some(session) = self.sessions.lock().await.attachments.remove(session_id) { + if let Some(cancellation) = session.cancellation { + cancellation.cancel(); + } + } + } + + pub async fn delete_session(&self, session_id: &str) -> Result<(), Error> { + self.close_session(session_id).await; + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + match manager.try_delete_session(session_id) { + Ok(true) => Ok(()), + Ok(false) => Err(Error::invalid_params().data("unknown session")), + Err(_) => Err(internal_error()), + } + } + + pub async fn cancel_session(&self, session_id: &str) { + if let Some(session) = self.sessions.lock().await.attachments.get(session_id) { + if let Some(cancellation) = &session.cancellation { + cancellation.cancel(); + } + } + } + + async fn attach_persisted_session( + &self, + session_id: &str, + cwd: PathBuf, + ) -> Result<(AcpSession, Vec), Error> { + // Fail promptly during close cleanup, before model/config discovery. Recheck + // below after awaits and hold the lock through the actual attachment. + self.sessions.lock().await.ensure_attachable(session_id)?; + let cwd = workspace_path(&cwd)?; + // Invalid load/resume requests must not wait on network model discovery. + { + let manager = self.session_manager.lock().map_err(|_| internal_error())?; + let stored = manager + .get_session_ref(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + if stored.workspace_path != cwd.to_string_lossy() { + return Err(Error::invalid_params().data("session does not belong to cwd")); + } + } + let config = crate::config::ConfigLoader::load_for(&cwd).map_err(|_| internal_error())?; + crate::skill::init_skill_store(&config.xdg_config_home, &config.project_root); + let models = crate::model::catalog::selectable_models(&config, None) + .await + .map_err(|_| internal_error())?; + + let mut sessions = self.sessions.lock().await; + sessions.ensure_attachable(session_id)?; + let messages = { + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + let stored = manager + .get_session(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + if stored.workspace_path != cwd.to_string_lossy() { + return Err(Error::invalid_params().data("session does not belong to cwd")); + } + stored.messages.clone() + }; + + let (provider, model) = messages + .iter() + .rev() + .find_map(|message| match (&message.provider, &message.model) { + (Some(provider), Some(model)) => Some((provider.clone(), model.clone())), + _ => None, + }) + .unwrap_or_else(|| resolve_model(&config)); + let agent = messages + .iter() + .rev() + .find_map(|message| message.agent_mode.clone()) + .unwrap_or_else(|| { + config + .merged_config + .default_agent + .clone() + .unwrap_or_else(|| { + config + .merged_config + .agent_registry + .default_agent() + .to_string() + }) + }); + let reasoning = model_reasoning(&config, &models, &provider, &model); + let reasoning_selection = + reasoning.unwrap_or(crate::model::reasoning::ReasoningEffort::None); + let context_window = model_context_window(&config, &models, &provider, &model); + let skills = crate::skill::SkillStore::load(&config.xdg_config_home, &config.project_root); + let tool_permissions = configured_tool_permissions(&cwd, &config); + let session = AcpSession { + cwd, + config, + tool_permissions, + skills, + models, + provider, + model, + agent, + reasoning_selection, + reasoning, + context_window, + cancellation: None, + }; + sessions + .attachments + .insert(session_id.to_string(), session.clone()); + Ok((session, messages)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::Path; + + #[test] + fn session_info_includes_hierarchy_metadata() { + let now = std::time::SystemTime::now(); + let info = acp_session_info( + crate::session::manager::SessionInfo { + id: "child".to_string(), + parent_id: Some("parent".to_string()), + title: "Child".to_string(), + created_at: now, + updated_at: now, + message_count: 0, + workspace_id: 1, + workspace_path: "/tmp".to_string(), + workspace_name: "tmp".to_string(), + workspace_sort_order: 0, + status: crate::session::types::SessionStatus::Idle, + pinned_at: None, + archived_at: None, + }, + Some("root".to_string()), + ); + let meta = info.meta.expect("session metadata"); + assert_eq!(meta["crabcode"]["parentSessionId"], "parent"); + assert_eq!(meta["crabcode"]["rootSessionId"], "root"); + } + + #[test] + fn session_list_cursor_pages_all_sessions() { + assert_eq!( + session_page(None, 250).unwrap(), + (0, 100, Some("100".to_string())) + ); + assert_eq!( + session_page(Some("100"), 250).unwrap(), + (100, 200, Some("200".to_string())) + ); + assert_eq!(session_page(Some("200"), 250).unwrap(), (200, 250, None)); + assert!(session_page(Some("invalid"), 250).is_err()); + } + + #[test] + fn rejects_non_absolute_workspaces() { + assert!(workspace_path(Path::new("relative")).is_err()); + } +} diff --git a/src/acp/service/stream.rs b/src/acp/service/stream.rs new file mode 100644 index 0000000..2a6d68e --- /dev/null +++ b/src/acp/service/stream.rs @@ -0,0 +1,550 @@ +use super::cancellable; +use super::internal_error; +use super::internal_error_with; +use super::notifications::send_text; +use super::notifications::send_tool_call; +use super::notifications::send_tool_result; +use super::permissions::request_permission; +use super::questions::request_questions; +use super::questions::skipped_question_answers; +use super::terminals::bridge_terminal_session; +use super::usage::cumulative_cost; +use super::usage::estimate_session_usage_cost; +use super::usage::send_live_usage; +use super::usage::send_usage; +use super::usage::AcpLiveUsage; +use super::AcpService; +use super::AcpSession; +use agent_client_protocol::schema::v1::PromptResponse; +use agent_client_protocol::schema::v1::StopReason; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use agent_client_protocol::Error; +use tokio_util::sync::CancellationToken; + +fn acp_stop_reason(reason: Option) -> StopReason { + match reason { + Some(crate::llm::TurnStopReason::MaxTokens) => StopReason::MaxTokens, + Some(crate::llm::TurnStopReason::MaxTurnRequests) => StopReason::MaxTurnRequests, + Some(crate::llm::TurnStopReason::Refusal) => StopReason::Refusal, + None => StopReason::EndTurn, + } +} + +pub(super) fn record_tool_call( + assistant: &mut crate::session::types::Message, + tool_call: &crate::llm::ToolCall, +) { + let args = serde_json::from_str(&tool_call.function.arguments) + .unwrap_or_else(|_| serde_json::Value::String(tool_call.function.arguments.clone())); + let provider_executed = matches!( + tool_call.function.name.as_str(), + "x_search" | "web_search" | "file_search" + ); + + if let Some(part) = assistant + .parts + .iter_mut() + .find(|part| part.part_type == "tool_call" && part.tool_id() == Some(tool_call.id.as_str())) + { + if let Some(obj) = part.data.as_object_mut() { + obj.insert( + "name".to_string(), + serde_json::Value::String(tool_call.function.name.clone()), + ); + obj.insert("args".to_string(), args); + if provider_executed { + obj.insert( + "provider_executed".to_string(), + serde_json::Value::Bool(true), + ); + } + } + return; + } + + assistant.add_tool_call_part(tool_call.id.clone(), tool_call.function.name.clone(), args); + if provider_executed { + if let Some(obj) = assistant + .parts + .last_mut() + .and_then(|part| part.data.as_object_mut()) + { + obj.insert( + "provider_executed".to_string(), + serde_json::Value::Bool(true), + ); + } + } +} + +pub(super) fn record_tool_result( + assistant: &mut crate::session::types::Message, + result: &crate::llm::ToolCallResult, +) { + let mut data = assistant + .tool_call_part_data(&result.tool_call_id) + .cloned() + .unwrap_or_else(|| serde_json::json!({})); + data["id"] = serde_json::Value::String(result.tool_call_id.clone()); + data["name"] = serde_json::Value::String(result.name.clone()); + if matches!( + result.name.as_str(), + "x_search" | "web_search" | "file_search" + ) { + data["provider_executed"] = serde_json::Value::Bool(true); + } + + match serde_json::from_str::(&result.content) { + Ok(serde_json::Value::Object(payload)) => { + if let Some(data) = data.as_object_mut() { + for (key, value) in payload { + data.insert(key, value); + } + } + } + _ => { + data["status"] = serde_json::Value::String("ok".to_string()); + data["output_preview"] = serde_json::Value::String(result.content.clone()); + } + } + assistant.add_or_update_tool_result_part(data); +} + +impl AcpService { + // Keep finalization outside every cancellable await, including client requests. + pub(super) async fn receive_prompt_stream( + &self, + session_id: &str, + session: &AcpSession, + connection: &ConnectionTo, + cancellation: &CancellationToken, + mut receiver: tokio::sync::mpsc::UnboundedReceiver, + base_context_tokens: usize, + base_prompt_tokens: usize, + base_cost: f64, + ) -> Result { + let mut assistant = crate::session::types::Message::incomplete(""); + let message_id = assistant.id.clone(); + assistant.provider = Some(session.provider.clone()); + assistant.model = Some(session.model.clone()); + assistant.agent_mode = Some(session.agent.clone()); + let mut failed = None; + let mut cancelled = false; + let mut turn_stop_reason = None; + let mut live_usage = AcpLiveUsage::default(); + + // Notification failures must take the same persistence path as cancellation. + let stream_result: Result<(), Error> = async { + while let Some(chunk) = cancellable(cancellation, receiver.recv()).await.flatten() { + match chunk { + crate::llm::ChunkMessage::Text(text) => { + assistant.append(&text); + send_text(connection, session_id, &message_id, text, false)?; + send_live_usage( + connection, + session_id, + session, + base_context_tokens, + base_cost, + &assistant, + &live_usage, + )?; + } + crate::llm::ChunkMessage::Reasoning(text) => { + assistant.append_reasoning(&text); + send_text(connection, session_id, &message_id, text, true)?; + send_live_usage( + connection, + session_id, + session, + base_context_tokens, + base_cost, + &assistant, + &live_usage, + )?; + } + crate::llm::ChunkMessage::ToolCalls(tool_calls) => { + for tool_call in &tool_calls { + record_tool_call(&mut assistant, tool_call); + } + for tool_call in tool_calls { + send_tool_call(connection, session_id, tool_call, &session.cwd)?; + } + send_live_usage( + connection, + session_id, + session, + base_context_tokens, + base_cost, + &assistant, + &live_usage, + )?; + } + crate::llm::ChunkMessage::ToolResult(result) => { + record_tool_result(&mut assistant, &result); + send_tool_result(connection, session_id, result, &session.cwd)?; + send_live_usage( + connection, + session_id, + session, + base_context_tokens, + base_cost, + &assistant, + &live_usage, + )?; + } + crate::llm::ChunkMessage::Metrics { + token_count, + duration_ms, + usage, + cost, + } => { + assistant.token_count = Some(token_count); + assistant.duration_ms = Some(duration_ms); + if let Some(usage) = usage { + assistant.apply_usage(usage, cost); + live_usage.cumulative = usage; + } + live_usage.cost = cost.unwrap_or(live_usage.cost); + send_usage( + connection, + session_id, + session, + live_usage + .context_tokens(base_context_tokens, &assistant) + .max(base_context_tokens.saturating_add(token_count)), + cumulative_cost(base_cost, live_usage.cost), + usage, + )?; + } + crate::llm::ChunkMessage::Usage(usage) => { + live_usage.observe_provider_usage(&assistant, usage); + live_usage.cost += estimate_session_usage_cost(session, &usage); + assistant + .parts + .push(crate::session::types::MessagePart::usage( + usage.input, + usage.output, + usage.cache_read, + usage.cache_write, + estimate_session_usage_cost(session, &usage), + )); + if let Some(part) = assistant.parts.last_mut() { + part.data["acp_prompt_tokens"] = serde_json::json!(base_prompt_tokens); + } + if usage.output > 0 { + assistant.output_tokens = Some( + assistant + .output_tokens + .unwrap_or(0) + .saturating_add(usage.output as usize), + ); + } + send_live_usage( + connection, + session_id, + session, + base_context_tokens, + base_cost, + &assistant, + &live_usage, + )?; + } + crate::llm::ChunkMessage::Cancelled => { + cancelled = true; + break; + } + crate::llm::ChunkMessage::Failed(error) => failed = Some(error), + crate::llm::ChunkMessage::TurnStopReason(reason) => { + turn_stop_reason = Some(reason) + } + crate::llm::ChunkMessage::PermissionRequest(prompt) => { + let response = cancellable( + cancellation, + request_permission(connection, session_id, &prompt), + ) + .await + .unwrap_or(crate::tools::PermissionResponse::Deny); + let _ = prompt.response_tx.send(response); + } + crate::llm::ChunkMessage::QuestionRequest { + tool_call_id, + questions, + response_tx, + } => { + let response = if self.supports_form_elicitation() { + request_questions( + connection, + session_id, + tool_call_id.as_deref(), + &questions, + cancellation, + ) + .await + } else { + skipped_question_answers(&questions) + }; + let _ = response_tx.send(response); + } + crate::llm::ChunkMessage::TerminalSessionRequest(request) => { + if self.supports_terminals() { + bridge_terminal_session( + connection, + session_id, + &session.cwd, + request, + cancellation, + ) + .await; + } else { + let _ = request + .control_tx + .send(crate::tools::TerminalSessionControl::Stop); + } + } + crate::llm::ChunkMessage::End => break, + _ => {} + } + } + + Ok(()) + } + .await; + if stream_result.is_err() && failed.is_none() { + failed = Some("ACP stream notification failed".to_string()); + } + assistant.is_complete = true; + assistant.was_interrupted = cancelled || cancellation.is_cancelled(); + { + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + let persist_result = manager.add_message_to_session(session_id, &assistant); + let status = if assistant.was_interrupted { + crate::session::types::SessionStatus::Interrupted + } else if failed.is_some() { + crate::session::types::SessionStatus::Failed + } else { + crate::session::types::SessionStatus::Idle + }; + manager + .set_session_status(session_id, status, failed.as_deref()) + .map_err(|_| internal_error())?; + persist_result.map_err(|_| internal_error())?; + } + if assistant.was_interrupted { + return Ok(PromptResponse::new(StopReason::Cancelled)); + } + stream_result?; + if let Some(error) = failed { + return Err(internal_error_with(&error)); + } + Ok(PromptResponse::new(acp_stop_reason(turn_stop_reason))) + } +} + +#[cfg(test)] +mod tests { + use super::super::test_support::empty_config; + use super::super::test_support::session_with_config; + use super::super::AcpService; + use super::super::AcpSessions; + use super::*; + use crate::session::manager::SessionManager; + use agent_client_protocol::schema::v1::RequestPermissionRequest; + use agent_client_protocol::schema::v1::SessionNotification; + use agent_client_protocol::schema::v1::SessionUpdate; + use agent_client_protocol::schema::v1::StopReason; + use agent_client_protocol::Client; + use std::sync::Arc; + use std::sync::Mutex; + use tokio::sync::Mutex as AsyncMutex; + + #[test] + fn maps_typed_turn_stop_reasons_to_acp() { + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTokens)), + StopReason::MaxTokens + ); + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::Refusal)), + StopReason::Refusal + ); + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), + StopReason::MaxTurnRequests + ); + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), + StopReason::MaxTurnRequests + ); + assert_eq!(acp_stop_reason(None), StopReason::EndTurn); + } + + #[tokio::test(flavor = "current_thread")] + async fn cancelled_stream_persists_partial_transcript_and_interrupted_status() { + cancelled_stream_finalizes(false, false).await; + } + + #[tokio::test(flavor = "current_thread")] + async fn cancelled_permission_persists_partial_transcript_and_interrupted_status() { + cancelled_stream_finalizes(true, false).await; + } + + #[tokio::test(flavor = "current_thread")] + async fn close_reopen_waits_for_old_prompt_finalization() { + cancelled_stream_finalizes(false, true).await; + } + + async fn cancelled_stream_finalizes(pending_permission: bool, close_reopen: bool) { + use agent_client_protocol::schema::v1::RequestPermissionResponse; + use agent_client_protocol::{Agent, Responder}; + use std::time::Duration; + + tokio::task::LocalSet::new().run_until(async { + let workspace = tempfile::tempdir().unwrap(); + let mut manager = SessionManager::new() + .with_history_for_workspace(workspace.path()).unwrap(); + let session_id = manager.create_session(Some("ACP cancellation regression".into())); + let db_id = manager.get_db_id(&session_id).unwrap(); + manager.set_session_status( + &session_id, crate::session::types::SessionStatus::Streaming, None, + ).unwrap(); + let service = AcpService { + sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), + session_manager: Arc::new(Mutex::new(manager)), + client_capabilities: Arc::new(Mutex::new(Default::default())), + }; + let session = session_with_config(empty_config()); + service.sessions.lock().await.attachments.insert(session_id.clone(), session.clone()); + let cancellation = service.sessions.lock().await.begin_prompt(&session_id).unwrap(); + let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); + let (permission_tx, mut permission_rx) = tokio::sync::mpsc::unbounded_channel(); + let client = Client.builder().on_receive_request( + async move |_request: RequestPermissionRequest, + responder: Responder, _connection| { + permission_tx.send(responder).unwrap(); + Ok(()) + }, + agent_client_protocol::on_receive_request!(), + ).on_receive_notification( + async move |notification: SessionNotification, _connection| { + if matches!(notification.update, SessionUpdate::AgentMessageChunk(_)) { + let _ = observed_tx.send(()); + } + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ); + Agent.builder().connect_with(client, async move |connection| { + let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); + sender.send(crate::llm::ChunkMessage::Text("partial answer".into())).unwrap(); + sender.send(crate::llm::ChunkMessage::Reasoning("partial reasoning".into())).unwrap(); + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + if pending_permission { + sender.send(crate::llm::ChunkMessage::PermissionRequest(crate::tools::PermissionPrompt { + tool_call_id: Some("pending-tool".into()), + tool_id: "bash".into(), + action: crate::tools::PermissionAction::Write, + permission: "bash".into(), + patterns: Vec::new(), + target: None, + command: Some("mock-command".into()), + workdir: None, + workspace: workspace.path().to_string_lossy().into_owned(), + reason: "regression test".into(), + raw_input: serde_json::json!({"command": "mock-command"}), + response_tx, + })).unwrap(); + } + let stream = service.receive_prompt_stream( + &session_id, &session, &connection, &cancellation, receiver, 0, 0, 0.0, + ); + tokio::pin!(stream); + tokio::time::timeout(Duration::from_secs(2), async { + tokio::select! { + _ = observed_rx.recv() => {}, + result = &mut stream => panic!("stream ended before cancellation: {result:?}"), + } + }).await.unwrap(); + let _pending_responder = if pending_permission { + Some(tokio::time::timeout(Duration::from_secs(2), async { + tokio::select! { + responder = permission_rx.recv() => responder.unwrap(), + result = &mut stream => panic!("stream ended before permission: {result:?}"), + } + }).await.unwrap()) + } else { + None + }; + if close_reopen { + service.close_session(&session_id).await; + assert!(cancellation.is_cancelled()); + assert!(!service.sessions.lock().await.attachments.contains_key(&session_id)); + assert!(service.available_commands(&session_id).await.is_err()); + assert!(service.set_mode(&session_id, "Build").await.is_err()); + assert!(service.set_model(&session_id, "unused").await.is_err()); + assert!(service.set_reasoning_effort(&session_id, "none").await.is_err()); + // The same guard serves load and resume. Resume must fail before + // any provider discovery, while the old stream still owns writes. + assert!(service.resume_session(session_id.clone(), workspace.path().into()).await.is_err()); + assert!(service.load_session(session_id.clone(), workspace.path().into(), connection.clone()).await.is_err()); + assert!(service.sessions.lock().await.begin_prompt(&session_id).is_err()); + } else { + cancellation.cancel(); + } + let response = tokio::time::timeout(Duration::from_secs(1), &mut stream) + .await.unwrap().unwrap(); + assert_eq!(response.stop_reason, StopReason::Cancelled); + if pending_permission { + assert!(matches!(response_rx.await.unwrap(), crate::tools::PermissionResponse::Deny)); + } + let manager = service.session_manager.lock().unwrap(); + let stored = manager.get_session_ref(&session_id).unwrap(); + assert_eq!(stored.status, crate::session::types::SessionStatus::Interrupted); + assert_eq!(stored.messages.len(), 1); + let assistant = &stored.messages[0]; + assert_eq!(assistant.content, "partial answer"); + assert!(assistant.is_complete); + assert!(assistant.was_interrupted); + drop(manager); + // A fresh DAO verifies the real database, not just the in-memory message. + let dao = crate::persistence::HistoryDAO::new_for_workspace(workspace.path()).unwrap(); + assert_eq!(dao.get_session(db_id).unwrap().unwrap().status, "interrupted"); + let messages = dao.get_messages(db_id).unwrap(); + assert_eq!(messages.len(), 1); + assert!(messages[0].parts.iter().any(|part| + part.part_type == "text" && part.data["text"] == "partial answer" + )); + drop(dao); + if close_reopen { + // Stream persistence has ended, but wrapper finalization has not: + // reopen must stay blocked through that second status write too. + assert!(service.resume_session(session_id.clone(), workspace.path().into()).await.is_err()); + service.session_manager.lock().unwrap().set_session_status( + &session_id, crate::session::types::SessionStatus::Interrupted, None, + ).unwrap(); + service.sessions.lock().await.finish_prompt(&session_id); + let mut sessions = service.sessions.lock().await; + sessions.ensure_attachable(&session_id).unwrap(); + sessions.attachments.insert(session_id.clone(), session.clone()); + let next = sessions.begin_prompt(&session_id).unwrap(); + drop(sessions); + service.session_manager.lock().unwrap().set_session_status( + &session_id, crate::session::types::SessionStatus::Streaming, None, + ).unwrap(); + assert!(!next.is_cancelled()); + assert!(service.sessions.lock().await.active_prompts.contains(&session_id)); + let dao = crate::persistence::HistoryDAO::new_for_workspace(workspace.path()).unwrap(); + assert_eq!(dao.get_session(db_id).unwrap().unwrap().status, "streaming"); + service.cancel_session(&session_id).await; + assert!(next.is_cancelled(), "the new turn retains its cancellation token"); + service.sessions.lock().await.finish_prompt(&session_id); + } else { + service.sessions.lock().await.finish_prompt(&session_id); + } + service.session_manager.lock().unwrap().try_delete_session(&session_id).unwrap(); + Ok(()) + }).await.unwrap(); + }).await; + } +} diff --git a/src/acp/service/terminals.rs b/src/acp/service/terminals.rs new file mode 100644 index 0000000..c9c8418 --- /dev/null +++ b/src/acp/service/terminals.rs @@ -0,0 +1,371 @@ +use super::cancellable; +use super::notifications::absolute_tool_path; +use agent_client_protocol::schema::v1::CreateTerminalRequest; +use agent_client_protocol::schema::v1::KillTerminalRequest; +use agent_client_protocol::schema::v1::ReleaseTerminalRequest; +use agent_client_protocol::schema::v1::SessionNotification; +use agent_client_protocol::schema::v1::SessionUpdate; +use agent_client_protocol::schema::v1::Terminal; +use agent_client_protocol::schema::v1::TerminalOutputRequest; +use agent_client_protocol::schema::v1::ToolCallContent; +use agent_client_protocol::schema::v1::ToolCallUpdate; +use agent_client_protocol::schema::v1::ToolCallUpdateFields; +use agent_client_protocol::schema::v1::WaitForTerminalExitRequest; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use std::path::Path; +use std::path::PathBuf; +use tokio_util::sync::CancellationToken; + +pub(super) async fn bridge_terminal_session( + connection: &ConnectionTo, + session_id: &str, + session_cwd: &Path, + request: crate::tools::TerminalSessionRequest, + cancellation: &CancellationToken, +) { + // Connection-owned task retains the create response even when the prompt stops. + // Do not tokio::spawn SDK block_task futures: they need the connection task context. + let worker_connection = connection.clone(); + let worker_session_id = session_id.to_string(); + let worker_cwd = session_cwd.to_path_buf(); + let worker_cancellation = cancellation.clone(); + let control_tx = request.control_tx.clone(); + let (done_tx, done_rx) = tokio::sync::oneshot::channel(); + if connection + .spawn(async move { + terminal_session_worker( + &worker_connection, + &worker_session_id, + &worker_cwd, + request, + &worker_cancellation, + ) + .await; + let _ = done_tx.send(()); + Ok(()) + }) + .is_err() + { + let _ = control_tx.send(crate::tools::TerminalSessionControl::Stop); + return; + } + if cancellable(cancellation, done_rx).await.is_none() { + let _ = control_tx.send(crate::tools::TerminalSessionControl::Stop); + } +} + +async fn cleanup_terminal( + connection: &ConnectionTo, + session_id: &str, + terminal_id: agent_client_protocol::schema::v1::TerminalId, +) { + // Issue both requests even if a client never acknowledges kill. + let kill = connection + .send_request(KillTerminalRequest::new( + session_id.to_string(), + terminal_id.clone(), + )) + .block_task(); + let release = connection + .send_request(ReleaseTerminalRequest::new( + session_id.to_string(), + terminal_id, + )) + .block_task(); + let _ = tokio::join!(kill, release); +} + +async fn terminal_session_worker( + connection: &ConnectionTo, + session_id: &str, + session_cwd: &Path, + request: crate::tools::TerminalSessionRequest, + cancellation: &CancellationToken, +) { + let start = request.start; + let control_tx = request.control_tx; + let cwd = start + .workdir + .as_deref() + .map(PathBuf::from) + .map(|path| absolute_tool_path(&path.to_string_lossy(), session_cwd)) + .unwrap_or_else(|| session_cwd.to_path_buf()); + let create = CreateTerminalRequest::new(session_id.to_string(), "bash") + .args(vec!["-c".to_string(), start.command.clone()]) + .cwd(cwd) + .output_byte_limit(crate::tools::terminal_session::MAX_TRANSCRIPT_BYTES as u64); + let terminal_id = match connection.send_request(create).block_task().await { + Ok(response) => response.terminal_id, + Err(error) => { + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( + format!("ACP client could not create terminal: {error}"), + )); + return; + } + }; + + if cancellation.is_cancelled() { + cleanup_terminal(connection, session_id, terminal_id).await; + return; + } + + let terminal_update = SessionUpdate::ToolCallUpdate(ToolCallUpdate::new( + start.tool_call_id.clone(), + ToolCallUpdateFields::new().content(vec![ToolCallContent::Terminal(Terminal::new( + terminal_id.clone(), + ))]), + )); + if connection + .send_notification(SessionNotification::new( + session_id.to_string(), + terminal_update, + )) + .is_err() + { + cleanup_terminal(connection, session_id, terminal_id).await; + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( + "ACP client could not embed terminal".to_string(), + )); + return; + } + + let wait = connection + .send_request(WaitForTerminalExitRequest::new( + session_id.to_string(), + terminal_id.clone(), + )) + .block_task(); + tokio::pin!(wait); + let (stopped_by_user, exit_code, wait_error) = tokio::select! { + _ = cancellation.cancelled() => { + cleanup_terminal(connection, session_id, terminal_id).await; + return; + } + response = &mut wait => match response { + Ok(response) => ( + false, + response.exit_status.exit_code.and_then(|code| i32::try_from(code).ok()), + None, + ), + Err(error) => (false, None, Some(error.to_string())), + } + }; + + let output = cancellable( + cancellation, + connection + .send_request(TerminalOutputRequest::new( + session_id.to_string(), + terminal_id.clone(), + )) + .block_task(), + ) + .await; + let Some(output) = output else { + cleanup_terminal(connection, session_id, terminal_id).await; + return; + }; + let _ = connection + .send_request(ReleaseTerminalRequest::new( + session_id.to_string(), + terminal_id, + )) + .block_task() + .await; + + match (wait_error, output) { + (None, Ok(output)) => { + let result = crate::tools::terminal_session::external_terminal_result( + &start, + &output.output, + output.truncated, + exit_code, + stopped_by_user, + ); + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalResult(result)); + } + (Some(error), _) => { + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( + format!("ACP terminal wait failed: {error}"), + )); + } + (None, Err(error)) => { + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( + format!("ACP terminal output failed: {error}"), + )); + } + } +} + +#[cfg(test)] +mod tests { + use super::super::AcpService; + use super::*; + use agent_client_protocol::schema::v1::CreateTerminalRequest; + use agent_client_protocol::schema::v1::KillTerminalRequest; + use agent_client_protocol::schema::v1::ReleaseTerminalRequest; + use agent_client_protocol::Client; + use agent_client_protocol::ConnectionTo; + use std::path::Path; + use tokio_util::sync::CancellationToken; + + #[test] + fn terminal_support_tracks_client_capability() { + let service = AcpService::new(Path::new("/tmp")).unwrap(); + assert!(!service.supports_terminals()); + + service.set_client_capabilities( + agent_client_protocol::schema::v1::ClientCapabilities::new().terminal(true), + ); + assert!(service.supports_terminals()); + } + + async fn delayed_terminal_creation_cleanup(acknowledge_kill: bool) { + use agent_client_protocol::schema::v1::{ + CreateTerminalResponse, KillTerminalResponse, ReleaseTerminalResponse, + }; + use agent_client_protocol::{Agent, Responder}; + use std::time::Duration; + + tokio::task::LocalSet::new() + .run_until(async { + let cancellation = CancellationToken::new(); + let (create_tx, mut create_rx) = tokio::sync::mpsc::unbounded_channel(); + let (cleanup_tx, mut cleanup_rx) = tokio::sync::mpsc::unbounded_channel(); + let kill_tx = cleanup_tx.clone(); + let (pending_kill_tx, mut pending_kill_rx) = + tokio::sync::mpsc::unbounded_channel(); + let client = Client + .builder() + .on_receive_request( + async move |request: CreateTerminalRequest, + responder: Responder, + _cx: ConnectionTo| { + assert_eq!(request.session_id.to_string(), "cancelled-session"); + assert_eq!(request.command, "bash"); + // Retain the response until the prompt-side bridge has returned. + create_tx.send(responder).unwrap(); + Ok(()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: KillTerminalRequest, + responder: Responder, + _cx: ConnectionTo| { + kill_tx + .send(("kill", request.session_id, request.terminal_id)) + .unwrap(); + if acknowledge_kill { + responder.respond(KillTerminalResponse::new()) + } else { + // Handlers run in the client's dispatch loop: waiting here + // would prevent it from dispatching terminal/release at all. + // Retain the responder without replying, but let dispatch + // continue so release must arrive with kill still outstanding. + pending_kill_tx.send(responder).unwrap(); + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: ReleaseTerminalRequest, + responder: Responder, + _cx: ConnectionTo| { + cleanup_tx + .send(("release", request.session_id, request.terminal_id)) + .unwrap(); + responder.respond(ReleaseTerminalResponse::new()) + }, + agent_client_protocol::on_receive_request!(), + ); + + let result = Agent + .builder() + .connect_with(client, async move |connection| { + let (control_tx, mut control_rx) = + tokio::sync::mpsc::unbounded_channel(); + let request = crate::tools::TerminalSessionRequest { + start: crate::tools::TerminalSessionStart { + session_id: "cancelled-session".to_string(), + tool_call_id: "terminal-call".to_string(), + command: "mock-command-never-executed".to_string(), + description: "Delayed mock terminal".to_string(), + workdir: None, + cols: 80, + rows: 24, + job_id: None, + }, + control_tx, + }; + let bridge = bridge_terminal_session( + &connection, + "cancelled-session", + Path::new("/tmp"), + request, + &cancellation, + ); + tokio::pin!(bridge); + let responder = tokio::time::timeout(Duration::from_secs(2), async { + tokio::select! { + responder = create_rx.recv() => responder.unwrap(), + _ = &mut bridge => panic!("bridge returned before creation or cancellation"), + } + }) + .await + .expect("mock client did not receive terminal/create"); + + cancellation.cancel(); + tokio::time::timeout(Duration::from_secs(1), &mut bridge) + .await + .expect("prompt-side bridge waited for delayed terminal creation"); + assert!(matches!( + control_rx.try_recv(), + Ok(crate::tools::TerminalSessionControl::Stop) + )); + assert!(cleanup_rx.try_recv().is_err()); + + // The terminal is only allocated after cancellation has returned. + responder.respond(CreateTerminalResponse::new("late-terminal"))?; + let mut cleanup = Vec::new(); + for _ in 0..2 { + let (method, session_id, terminal_id) = + tokio::time::timeout(Duration::from_secs(2), cleanup_rx.recv()) + .await + .expect("late terminal did not receive both cleanup requests") + .expect("mock cleanup channel closed"); + assert_eq!(session_id.to_string(), "cancelled-session"); + assert_eq!(terminal_id.to_string(), "late-terminal"); + cleanup.push(method); + } + cleanup.sort_unstable(); + assert_eq!(cleanup, ["kill", "release"]); + if !acknowledge_kill { + // Only drop the unanswered kill responder after observing release. + // Dropping it earlier would send an error response and accidentally + // allow cleanup that waits for kill acknowledgement to pass. + let pending_kill = pending_kill_rx + .try_recv() + .expect("kill must remain unanswered until release arrives"); + drop(pending_kill); + } + Ok(()) + }) + .await; + result.expect("mock ACP connection failed"); + }) + .await; + } + + #[tokio::test(flavor = "current_thread")] + async fn cancelled_terminal_creation_cleans_up_delayed_response() { + delayed_terminal_creation_cleanup(true).await; + } + + #[tokio::test(flavor = "current_thread")] + async fn cancelled_terminal_creation_releases_without_kill_acknowledgement() { + delayed_terminal_creation_cleanup(false).await; + } +} diff --git a/src/acp/service/test_support.rs b/src/acp/service/test_support.rs new file mode 100644 index 0000000..ddc666e --- /dev/null +++ b/src/acp/service/test_support.rs @@ -0,0 +1,102 @@ +use super::permissions::configured_tool_permissions; +use super::AcpSession; +use crate::config::configuration::LoadedConfig; +use std::path::Path; +use std::path::PathBuf; + +pub(super) fn model( + provider_id: &str, + provider_name: &str, + id: &str, + name: &str, +) -> crate::model::types::Model { + crate::model::types::Model { + id: id.to_string(), + name: name.to_string(), + family: String::new(), + provider_id: provider_id.to_string(), + provider_name: provider_name.to_string(), + attachment: false, + structured_output: false, + free: false, + local: false, + reasoning_options: Vec::new(), + context_window: None, + } +} + +pub(super) fn reasoning_model() -> crate::model::types::Model { + let mut model = model("openai", "OpenAI", "o3", "o3"); + model.reasoning_options = vec![crate::model::reasoning::ReasoningOption { + kind: "effort".to_string(), + values: vec!["low".to_string(), "medium".to_string(), "high".to_string()], + }]; + model +} + +pub(super) fn test_session() -> AcpSession { + AcpSession { + cwd: PathBuf::from("/tmp"), + config: crate::config::configuration::LoadedConfig { + merged_config: crate::config::configuration::MergedConfig::default(), + raw_merged: serde_json::Value::Null, + diagnostics: Default::default(), + inventory: Default::default(), + project_root: PathBuf::from("/tmp"), + cwd: PathBuf::from("/tmp"), + xdg_config_home: PathBuf::from("/tmp"), + }, + tool_permissions: crate::tools::ToolPermissions::new("/tmp"), + skills: crate::skill::SkillStore::load(Path::new("/tmp"), Path::new("/tmp")), + models: vec![model("example", "Example", "chat", "Chat")], + provider: "example".to_string(), + model: "chat".to_string(), + agent: "Build".to_string(), + reasoning_selection: crate::model::reasoning::ReasoningEffort::High, + reasoning: None, + context_window: None, + cancellation: None, + } +} + +pub(super) fn config_with_command(command: crate::command::custom::CustomCommand) -> LoadedConfig { + let mut merged_config = crate::config::configuration::MergedConfig::default(); + merged_config.commands.push(command); + config_with_merged(merged_config) +} + +pub(super) fn config_with_merged( + merged_config: crate::config::configuration::MergedConfig, +) -> LoadedConfig { + LoadedConfig { + merged_config, + raw_merged: serde_json::Value::Null, + diagnostics: Default::default(), + inventory: Default::default(), + project_root: PathBuf::from("/tmp"), + cwd: PathBuf::from("/tmp"), + xdg_config_home: PathBuf::from("/tmp"), + } +} + +pub(super) fn empty_config() -> LoadedConfig { + config_with_merged(crate::config::configuration::MergedConfig::default()) +} + +pub(super) fn session_with_config(config: LoadedConfig) -> AcpSession { + let skills = crate::skill::SkillStore::load(&config.xdg_config_home, &config.project_root); + AcpSession { + cwd: config.cwd.clone(), + tool_permissions: configured_tool_permissions(&config.cwd, &config), + config, + skills, + models: Vec::new(), + provider: String::new(), + model: String::new(), + agent: "Build".to_string(), + reasoning_selection: crate::model::reasoning::ReasoningEffort::None, + reasoning: None, + context_window: None, + cancellation: None, + } +} diff --git a/src/acp/service/usage.rs b/src/acp/service/usage.rs new file mode 100644 index 0000000..c37f6e4 --- /dev/null +++ b/src/acp/service/usage.rs @@ -0,0 +1,474 @@ +use super::internal_error; +use super::AcpSession; +use agent_client_protocol::schema::v1::Cost as AcpCost; +use agent_client_protocol::schema::v1::SessionNotification; +use agent_client_protocol::schema::v1::SessionUpdate; +use agent_client_protocol::schema::v1::UsageUpdate; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use agent_client_protocol::Error; + +pub(super) fn acp_session_cost(messages: &[crate::session::types::Message]) -> f64 { + messages + .iter() + .map(|message| message.cost.unwrap_or_else(|| message.usage_cost())) + .sum() +} + +pub(super) fn acp_context_tokens( + messages: &[crate::session::types::Message], + prompt_tokens: usize, +) -> usize { + let estimated = + crate::session::compaction::total_context_tokens(messages).saturating_add(prompt_tokens); + // Retained tail usage still describes the old context, as does the summary's billed usage. + let start = messages + .iter() + .rposition(crate::session::compaction::is_compaction_marker) + .map(|index| index + 1) + .unwrap_or(0); + for (index, message) in messages.iter().enumerate().skip(start).rev() { + if message.role != crate::session::types::MessageRole::Assistant { + continue; + } + // Each tool step bills its prompt again; cumulative usage is not a context size. + let provider_context = message + .parts + .iter() + .enumerate() + .filter_map(|(part_index, part)| { + let usage = part.recorded_usage()?; + let input = usage + .input + .saturating_add(usage.cache_read) + .saturating_add(usage.cache_write); + if input == 0 { + return None; + } + let mut tokens = + usize::try_from(input.saturating_add(usage.output)).unwrap_or(usize::MAX); + if let Some(previous_prompt) = part + .data + .get("acp_prompt_tokens") + .and_then(serde_json::Value::as_u64) + { + tokens = tokens + .saturating_sub(usize::try_from(previous_prompt).unwrap_or(usize::MAX)) + .saturating_add(prompt_tokens); + } + Some(tokens.saturating_add( + crate::session::compaction::assistant_parts_context_tokens( + &message.parts[part_index + 1..], + ), + )) + }) + .max() + .filter(|tokens| *tokens > 0); + if let Some(tokens) = provider_context { + let following = + crate::session::compaction::total_context_tokens(&messages[index + 1..]); + return estimated.max(tokens.saturating_add(following)); + } + } + estimated +} + +pub(super) fn estimate_session_usage_cost( + session: &AcpSession, + usage: &crate::aisdk::chunk::TokenUsage, +) -> f64 { + crate::model::discovery::Discovery::new_with_custom(Some( + session.config.merged_config.custom_providers.clone(), + )) + .ok() + .map(|discovery| { + discovery.estimate_usage_cost( + &session.provider, + &session.model, + usage.input, + usage.output, + usage.cache_read, + usage.cache_write, + ) + }) + .unwrap_or(0.0) +} + +pub(super) fn send_usage( + connection: &ConnectionTo, + session_id: &str, + session: &AcpSession, + used: usize, + cost: Option, + usage: Option, +) -> Result<(), Error> { + let update = SessionUpdate::UsageUpdate(usage_update(session, used, cost, usage)); + connection + .send_notification(SessionNotification::new(session_id.to_string(), update)) + .map_err(|_| internal_error()) +} + +pub(super) fn send_live_usage( + connection: &ConnectionTo, + session_id: &str, + session: &AcpSession, + base_context_tokens: usize, + base_cost: f64, + assistant: &crate::session::types::Message, + live_usage: &AcpLiveUsage, +) -> Result<(), Error> { + send_usage( + connection, + session_id, + session, + live_usage.context_tokens(base_context_tokens, assistant), + cumulative_cost(base_cost, live_usage.cost), + (!live_usage.cumulative.is_empty()).then_some(live_usage.cumulative), + ) +} + +#[derive(Debug, Default)] +pub(super) struct AcpLiveUsage { + pub(super) cumulative: crate::aisdk::chunk::TokenUsage, + provider_context_floor: usize, + pub(super) cost: f64, +} + +impl AcpLiveUsage { + pub(super) fn observe_provider_usage( + &mut self, + assistant: &crate::session::types::Message, + usage: crate::aisdk::chunk::TokenUsage, + ) { + self.cumulative = self.cumulative.saturating_add(usage); + let provider_input = usage + .input + .saturating_add(usage.cache_read) + .saturating_add(usage.cache_write); + let observed_assistant = live_assistant_context_tokens(assistant); + self.provider_context_floor = self + .provider_context_floor + .max(usize::try_from(provider_input).unwrap_or(usize::MAX)) + .max(observed_assistant); + } + + pub(super) fn context_tokens( + &self, + base_context_tokens: usize, + assistant: &crate::session::types::Message, + ) -> usize { + base_context_tokens + .saturating_add(live_assistant_context_tokens(assistant)) + .max(self.provider_context_floor) + } +} + +fn live_assistant_context_tokens(assistant: &crate::session::types::Message) -> usize { + let persisted = crate::session::compaction::message_context_tokens(assistant); + let reasoning = assistant + .reasoning + .as_deref() + .map(estimate_acp_tokens) + .unwrap_or(0); + persisted.saturating_add(reasoning) +} + +fn estimate_acp_tokens(content: &str) -> usize { + content.chars().count().saturating_add(3) / 4 +} + +pub(super) fn cumulative_cost(base_cost: f64, turn_cost: f64) -> Option { + let cost = base_cost + turn_cost; + (cost > 0.0).then_some(cost) +} + +fn usage_update( + session: &AcpSession, + used: usize, + cost: Option, + usage: Option, +) -> UsageUpdate { + let size = session + .context_window + .map(u64::from) + .unwrap_or_else(|| (used as u64).max(1)); + let mut meta = serde_json::Map::new(); + meta.insert( + "crabcode".to_string(), + serde_json::json!({ + "contextWindowKnown": session.context_window.is_some(), + "usage": usage.map(|usage| serde_json::json!({ + "inputTokens": usage.input, + "outputTokens": usage.output, + "cacheReadTokens": usage.cache_read, + "cacheWriteTokens": usage.cache_write, + })), + }), + ); + UsageUpdate::new(used as u64, size) + .cost(cost.map(|amount| AcpCost::new(amount, "USD"))) + .meta(meta) +} + +#[cfg(test)] +mod tests { + use super::super::stream::record_tool_call; + use super::super::stream::record_tool_result; + use super::super::test_support::session_with_config; + use super::super::test_support::test_session; + use super::*; + use crate::config::configuration::LoadedConfig; + use std::path::PathBuf; + + #[test] + fn acp_usage_update_includes_cumulative_usd_cost() { + let usage = crate::aisdk::chunk::TokenUsage { + input: 800, + output: 200, + cache_read: 500, + cache_write: 100, + }; + let update = usage_update(&test_session(), 1_000, Some(0.125), Some(usage)); + assert_eq!(update.cost.as_ref().map(|cost| cost.amount), Some(0.125)); + assert_eq!( + update.cost.as_ref().map(|cost| cost.currency.as_str()), + Some("USD") + ); + assert_eq!(update.size, 1_000); + let meta = update.meta.expect("usage metadata"); + assert_eq!(meta["crabcode"]["contextWindowKnown"], false); + assert_eq!(meta["crabcode"]["usage"]["inputTokens"], 800); + assert_eq!(meta["crabcode"]["usage"]["cacheWriteTokens"], 100); + } + + #[test] + fn acp_live_usage_advances_with_streamed_tools_and_provider_context() { + let mut assistant = crate::session::types::Message::incomplete("hello"); + record_tool_call( + &mut assistant, + &crate::llm::ToolCall { + id: "call_1".to_string(), + call_type: "function".to_string(), + function: crate::llm::FunctionCall { + name: "read".to_string(), + arguments: serde_json::json!({ "filePath": "src/main.rs" }).to_string(), + }, + }, + ); + record_tool_result( + &mut assistant, + &crate::llm::ToolCallResult { + tool_call_id: "call_1".to_string(), + role: "tool".to_string(), + name: "read".to_string(), + content: serde_json::json!({ + "status": "ok", + "output_preview": "x".repeat(4_000), + }) + .to_string(), + }, + ); + + let mut live = AcpLiveUsage::default(); + let before_provider_usage = live.context_tokens(18_000, &assistant); + assert!(before_provider_usage > 18_000); + assert!(assistant + .tool_call_part_data("call_1") + .and_then(|part| part.get("args")) + .is_some()); + + live.observe_provider_usage( + &assistant, + crate::aisdk::chunk::TokenUsage { + input: 70_000, + output: 500, + cache_read: 20_000, + cache_write: 0, + }, + ); + assert_eq!(live.context_tokens(18_000, &assistant), 90_000); + assert_eq!(live.cumulative.input, 70_000); + assert_eq!(live.cumulative.cache_read, 20_000); + } + + #[test] + fn acp_live_usage_keeps_largest_provider_context_across_tool_steps() { + let assistant = crate::session::types::Message::incomplete(""); + let mut live = AcpLiveUsage::default(); + live.observe_provider_usage( + &assistant, + crate::aisdk::chunk::TokenUsage { + input: 60_000, + output: 200, + cache_read: 20_000, + cache_write: 0, + }, + ); + live.observe_provider_usage( + &assistant, + crate::aisdk::chunk::TokenUsage { + input: 10_000, + output: 100, + cache_read: 40_000, + cache_write: 0, + }, + ); + + assert_eq!(live.context_tokens(18_000, &assistant), 80_000); + assert_eq!(live.cumulative.input, 70_000); + assert_eq!(live.cumulative.output, 300); + assert_eq!(live.cumulative.cache_read, 60_000); + } + + #[test] + fn acp_compaction_counts_provider_context_and_cached_input() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(60_000, 200, 20_000, 5_000, 0.0)); + let messages = vec![Message::user("question"), assistant]; + + assert_eq!(acp_context_tokens(&messages, 1_000), 85_200); + } + + #[test] + fn acp_compaction_does_not_sum_context_across_tool_requests() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(60_000, 200, 20_000, 0, 0.0)); + assistant + .parts + .push(MessagePart::usage(10_000, 100, 40_000, 0, 0.0)); + + assert_eq!( + acp_context_tokens(&[Message::user("question"), assistant], 0), + 80_200 + ); + } + + #[test] + fn acp_compaction_uses_latest_turn_usage_and_counts_pending_history() { + use crate::session::types::{Message, MessagePart}; + let mut old = Message::assistant("old answer"); + old.parts.push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); + let mut latest = Message::assistant("new answer"); + latest + .parts + .push(MessagePart::usage(50_000, 200, 0, 0, 0.0)); + let pending = Message::user("continue"); + let pending_tokens = crate::session::compaction::message_context_tokens(&pending); + let messages = vec![Message::user("question"), old, latest, pending]; + + assert_eq!(acp_context_tokens(&messages, 0), 50_200 + pending_tokens); + } + + #[test] + fn acp_compaction_counts_prompt_and_pending_tokens_without_provider_usage() { + use crate::config::configuration::CompactionConfig; + use crate::session::types::Message; + let messages = vec![Message::user("history ".repeat(2_000))]; + let history_tokens = crate::session::compaction::total_context_tokens(&messages); + let config = CompactionConfig::Settings { + auto: true, + prune: false, + reserved: Some(1_000), + }; + let context_window = Some(history_tokens as u32 + 3_000); + assert!(!crate::session::compaction::should_auto_compact( + &config, + acp_context_tokens(&messages, 1_500), + context_window, + None, + )); + assert!(crate::session::compaction::should_auto_compact( + &config, + acp_context_tokens(&messages, 1_500) + 500, + context_window, + None, + )); + } + + #[test] + fn acp_compaction_replaces_previous_prompt_estimate_in_provider_context() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + let mut usage = MessagePart::usage(80_000, 200, 0, 0, 0.0); + usage.data["acp_prompt_tokens"] = serde_json::json!(5_000); + assistant.parts.push(usage); + let messages = vec![Message::user("question"), assistant]; + assert_eq!(acp_context_tokens(&messages, 25_000), 100_200); + assert_eq!(acp_context_tokens(&messages, 1_000), 76_200); + } + + #[test] + fn acp_compaction_keeps_provider_context_when_new_input_usage_is_missing() { + use crate::session::types::{Message, MessagePart}; + for output in [0, 100] { + let mut previous = Message::assistant("previous answer"); + previous + .parts + .push(MessagePart::usage(85_000, 200, 0, 0, 0.0)); + let mut latest = Message::assistant("new answer"); + let mut usage = MessagePart::usage(0, output, 0, 0, 0.0); + usage.data["acp_prompt_tokens"] = serde_json::json!(1_000); + latest.parts.push(usage); + let following_tokens = crate::session::compaction::message_context_tokens(&latest); + assert_eq!( + acp_context_tokens(&[Message::user("question"), previous, latest], 2_000), + 85_200 + following_tokens + ); + } + } + + #[test] + fn acp_compaction_counts_tool_output_after_provider_usage() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(80_000, 200, 0, 0, 0.0)); + record_tool_result( + &mut assistant, + &crate::llm::ToolCallResult { + tool_call_id: "call_1".into(), + role: "tool".into(), + name: "read".into(), + content: serde_json::json!({"output_preview": "late tool output ".repeat(1_000)}) + .to_string(), + }, + ); + let mut following = Message::assistant(""); + following.parts = assistant.parts[2..].to_vec(); + let extra_tokens = crate::session::compaction::message_context_tokens(&following); + assert!(extra_tokens > 1_000); + assert_eq!( + acp_context_tokens(&[Message::user("question"), assistant], 0), + 80_200 + extra_tokens + ); + } + + #[test] + fn unknown_model_usage_cost_is_zero() { + let mut session = session_with_config(LoadedConfig { + merged_config: crate::config::configuration::MergedConfig::default(), + raw_merged: serde_json::Value::Null, + diagnostics: Default::default(), + inventory: Default::default(), + project_root: PathBuf::from("/tmp"), + cwd: PathBuf::from("/tmp"), + xdg_config_home: PathBuf::from("/tmp"), + }); + session.provider = "no-such-provider".to_string(); + session.model = "no-such-model".to_string(); + let usage = crate::aisdk::chunk::TokenUsage { + input: 1_000, + output: 1_000, + cache_read: 0, + cache_write: 0, + }; + assert_eq!(estimate_session_usage_cost(&session, &usage), 0.0); + } +} From 31891296621be4f0eb4710ff9c17b89e4b1e9600 Mon Sep 17 00:00:00 2001 From: Yanuar Date: Tue, 13 Oct 2026 12:20:44 +0700 Subject: [PATCH 4/5] refactor(acp): narrow service module responsibilities --- src/acp/service/attachment_storage.rs | 176 +++++++++ src/acp/service/attachments.rs | 171 +-------- src/acp/service/command_expansion.rs | 236 ++++++++++++ src/acp/service/commands.rs | 229 +----------- src/acp/service/compaction.rs | 380 +------------------ src/acp/service/compaction_summary.rs | 129 +++++++ src/acp/service/config.rs | 504 +------------------------- src/acp/service/context.rs | 197 ++++++++++ src/acp/service/mcp_config.rs | 112 ++++++ src/acp/service/mod.rs | 12 + src/acp/service/models.rs | 300 +++++++++++++++ src/acp/service/modes.rs | 106 ++++++ src/acp/service/notifications.rs | 489 +------------------------ src/acp/service/permission_preview.rs | 226 ++++++++++++ src/acp/service/permissions.rs | 223 +----------- src/acp/service/prompt.rs | 16 +- src/acp/service/replay.rs | 160 ++++++++ src/acp/service/sessions.rs | 63 +--- src/acp/service/stream.rs | 289 +-------------- src/acp/service/terminals.rs | 5 +- src/acp/service/tests/compaction.rs | 247 +++++++++++++ src/acp/service/tests/sessions.rs | 47 +++ src/acp/service/tests/stream.rs | 203 +++++++++++ src/acp/service/tool_content.rs | 242 +++++++++++++ src/acp/service/tool_locations.rs | 116 ++++++ src/acp/service/transcript.rs | 79 ++++ src/acp/service/usage.rs | 257 ++----------- 27 files changed, 2671 insertions(+), 2543 deletions(-) create mode 100644 src/acp/service/attachment_storage.rs create mode 100644 src/acp/service/command_expansion.rs create mode 100644 src/acp/service/compaction_summary.rs create mode 100644 src/acp/service/context.rs create mode 100644 src/acp/service/mcp_config.rs create mode 100644 src/acp/service/models.rs create mode 100644 src/acp/service/modes.rs create mode 100644 src/acp/service/permission_preview.rs create mode 100644 src/acp/service/replay.rs create mode 100644 src/acp/service/tests/compaction.rs create mode 100644 src/acp/service/tests/sessions.rs create mode 100644 src/acp/service/tests/stream.rs create mode 100644 src/acp/service/tool_content.rs create mode 100644 src/acp/service/tool_locations.rs create mode 100644 src/acp/service/transcript.rs diff --git a/src/acp/service/attachment_storage.rs b/src/acp/service/attachment_storage.rs new file mode 100644 index 0000000..bbcb0b5 --- /dev/null +++ b/src/acp/service/attachment_storage.rs @@ -0,0 +1,176 @@ +use super::internal_error; +use agent_client_protocol::Error; +use base64::Engine as _; +use std::path::Path; + +pub(super) fn write_prompt_audio( + session_id: &str, + audio: &agent_client_protocol::schema::v1::AudioContent, +) -> Result { + const MAX_AUDIO_BYTES: usize = 20 * 1024 * 1024; + let media_type = audio.mime_type.trim().to_ascii_lowercase(); + let extension = match media_type.as_str() { + "audio/wav" | "audio/x-wav" | "audio/wave" => "wav", + "audio/mpeg" | "audio/mp3" => "mp3", + _ => { + return Err(Error::invalid_params() + .data(format!("unsupported audio MIME type: {}", audio.mime_type))); + } + }; + let data = base64::engine::general_purpose::STANDARD + .decode(audio.data.trim()) + .map_err(|error| Error::invalid_params().data(format!("invalid audio data: {error}")))?; + if data.len() > MAX_AUDIO_BYTES { + return Err(Error::invalid_params().data("audio exceeds the 20 MiB size limit")); + } + let path = crate::persistence::attachments::write(session_id, extension, &data) + .map_err(|_| internal_error())?; + Ok(path.to_string_lossy().into_owned()) +} + +pub(super) struct ManagedAttachmentGuard { + paths: Vec, + committed: bool, +} + +pub(super) fn write_prompt_image( + session_id: &str, + image: &agent_client_protocol::schema::v1::ImageContent, +) -> Result { + const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024; + + let (mime_type, encoded_data) = prompt_image_payload(image)?; + let extension = match mime_type { + "image/png" => "png", + "image/jpeg" => "jpg", + "image/gif" => "gif", + "image/webp" => "webp", + mime_type => { + return Err( + Error::invalid_params().data(format!("unsupported image MIME type: {mime_type}")) + ); + } + }; + let data = base64::engine::general_purpose::STANDARD + .decode(encoded_data) + .map_err(|error| Error::invalid_params().data(format!("invalid image data: {error}")))?; + if data.len() > MAX_IMAGE_BYTES { + return Err(Error::invalid_params().data("image exceeds the 20 MiB size limit")); + } + + let path = crate::persistence::attachments::write(session_id, extension, &data) + .map_err(|_| internal_error())?; + Ok(path.to_string_lossy().into_owned()) +} + +fn prompt_image_payload( + image: &agent_client_protocol::schema::v1::ImageContent, +) -> Result<(&str, &str), Error> { + let data = image.data.trim(); + if data.is_empty() { + return Err(Error::invalid_params().data("image data is empty")); + } + + let Some(data_uri) = data.strip_prefix("data:") else { + return Ok((image.mime_type.as_str(), data)); + }; + let (metadata, payload) = data_uri + .split_once(',') + .ok_or_else(|| Error::invalid_params().data("invalid image data URI"))?; + let mut metadata = metadata.split(';'); + let mime_type = metadata + .next() + .filter(|mime_type| !mime_type.is_empty()) + .ok_or_else(|| Error::invalid_params().data("image data URI is missing a MIME type"))?; + if !metadata.any(|value| value.eq_ignore_ascii_case("base64")) { + return Err(Error::invalid_params().data("image data URI must be base64 encoded")); + } + let payload = payload.trim(); + if payload.is_empty() { + return Err(Error::invalid_params().data("image data is empty")); + } + + Ok((mime_type, payload)) +} + +impl ManagedAttachmentGuard { + pub(super) fn new(paths: Vec) -> Self { + Self { + paths, + committed: false, + } + } + + pub(super) fn commit(&mut self) { + self.committed = true; + } +} + +impl Drop for ManagedAttachmentGuard { + fn drop(&mut self) { + if !self.committed { + for path in &self.paths { + crate::persistence::attachments::remove_file(Path::new(path)); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::Path; + + #[test] + fn writes_supported_acp_image_to_managed_session_storage() { + let session_id = format!("acp-image-{}", cuid2::create_id()); + let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/png"); + let path = write_prompt_image(&session_id, &image).expect("image file"); + + assert!(Path::new(&path).starts_with(crate::persistence::attachments::root_dir())); + assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn writes_supported_acp_audio_to_managed_session_storage() { + let session_id = format!("acp-audio-{}", cuid2::create_id()); + let audio = agent_client_protocol::schema::v1::AudioContent::new("YXVkaW8=", "audio/wav"); + let path = write_prompt_audio(&session_id, &audio).expect("audio file"); + + assert!(path.ends_with(".wav")); + assert_eq!(std::fs::read(&path).unwrap(), b"audio"); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn writes_acp_clipboard_image_data_uri_to_managed_storage() { + let session_id = format!("acp-image-{}", cuid2::create_id()); + let image = agent_client_protocol::schema::v1::ImageContent::new( + "data:image/png;base64,aGk=", + "application/octet-stream", + ); + let path = write_prompt_image(&session_id, &image).expect("image file"); + + assert!(path.ends_with(".png")); + assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn rejects_non_base64_acp_image_data_uri() { + let image = agent_client_protocol::schema::v1::ImageContent::new( + "data:image/png,not-base64", + "image/png", + ); + + assert!(write_prompt_image("test", &image).is_err()); + } + + #[test] + fn rejects_unsupported_acp_image_mime_type() { + let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/tiff"); + + assert!(write_prompt_image("test", &image).is_err()); + } +} diff --git a/src/acp/service/attachments.rs b/src/acp/service/attachments.rs index 4c822d3..5a4f00b 100644 --- a/src/acp/service/attachments.rs +++ b/src/acp/service/attachments.rs @@ -1,37 +1,12 @@ -use super::internal_error; +use super::attachment_storage::write_prompt_audio; +use super::attachment_storage::write_prompt_image; use super::AcpSession; use crate::config::configuration::LoadedConfig; use agent_client_protocol::schema::v1::ContentBlock; use agent_client_protocol::schema::v1::EmbeddedResourceResource; use agent_client_protocol::Error; -use base64::Engine as _; use std::path::Path; -fn write_prompt_audio( - session_id: &str, - audio: &agent_client_protocol::schema::v1::AudioContent, -) -> Result { - const MAX_AUDIO_BYTES: usize = 20 * 1024 * 1024; - let media_type = audio.mime_type.trim().to_ascii_lowercase(); - let extension = match media_type.as_str() { - "audio/wav" | "audio/x-wav" | "audio/wave" => "wav", - "audio/mpeg" | "audio/mp3" => "mp3", - _ => { - return Err(Error::invalid_params() - .data(format!("unsupported audio MIME type: {}", audio.mime_type))); - } - }; - let data = base64::engine::general_purpose::STANDARD - .decode(audio.data.trim()) - .map_err(|error| Error::invalid_params().data(format!("invalid audio data: {error}")))?; - if data.len() > MAX_AUDIO_BYTES { - return Err(Error::invalid_params().data("audio exceeds the 20 MiB size limit")); - } - let path = crate::persistence::attachments::write(session_id, extension, &data) - .map_err(|_| internal_error())?; - Ok(path.to_string_lossy().into_owned()) -} - pub(super) fn model_supports_audio(config: &LoadedConfig, provider: &str, model: &str) -> bool { crate::model::discovery::Discovery::new_with_custom(Some( config.merged_config.custom_providers.clone(), @@ -40,34 +15,6 @@ pub(super) fn model_supports_audio(config: &LoadedConfig, provider: &str, model: .is_some_and(|discovery| discovery.model_supports_input_modality(provider, model, "audio")) } -pub(super) struct ManagedAttachmentGuard { - paths: Vec, - committed: bool, -} - -impl ManagedAttachmentGuard { - pub(super) fn new(paths: Vec) -> Self { - Self { - paths, - committed: false, - } - } - - pub(super) fn commit(&mut self) { - self.committed = true; - } -} - -impl Drop for ManagedAttachmentGuard { - fn drop(&mut self) { - if !self.committed { - for path in &self.paths { - crate::persistence::attachments::remove_file(Path::new(path)); - } - } - } -} - pub(super) fn prompt_content( parts: Vec, text_override: Option<&str>, @@ -173,73 +120,12 @@ fn prompt_text(parts: Vec) -> Result { Ok(text) } -fn write_prompt_image( - session_id: &str, - image: &agent_client_protocol::schema::v1::ImageContent, -) -> Result { - const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024; - - let (mime_type, encoded_data) = prompt_image_payload(image)?; - let extension = match mime_type { - "image/png" => "png", - "image/jpeg" => "jpg", - "image/gif" => "gif", - "image/webp" => "webp", - mime_type => { - return Err( - Error::invalid_params().data(format!("unsupported image MIME type: {mime_type}")) - ); - } - }; - let data = base64::engine::general_purpose::STANDARD - .decode(encoded_data) - .map_err(|error| Error::invalid_params().data(format!("invalid image data: {error}")))?; - if data.len() > MAX_IMAGE_BYTES { - return Err(Error::invalid_params().data("image exceeds the 20 MiB size limit")); - } - - let path = crate::persistence::attachments::write(session_id, extension, &data) - .map_err(|_| internal_error())?; - Ok(path.to_string_lossy().into_owned()) -} - -fn prompt_image_payload( - image: &agent_client_protocol::schema::v1::ImageContent, -) -> Result<(&str, &str), Error> { - let data = image.data.trim(); - if data.is_empty() { - return Err(Error::invalid_params().data("image data is empty")); - } - - let Some(data_uri) = data.strip_prefix("data:") else { - return Ok((image.mime_type.as_str(), data)); - }; - let (metadata, payload) = data_uri - .split_once(',') - .ok_or_else(|| Error::invalid_params().data("invalid image data URI"))?; - let mut metadata = metadata.split(';'); - let mime_type = metadata - .next() - .filter(|mime_type| !mime_type.is_empty()) - .ok_or_else(|| Error::invalid_params().data("image data URI is missing a MIME type"))?; - if !metadata.any(|value| value.eq_ignore_ascii_case("base64")) { - return Err(Error::invalid_params().data("image data URI must be base64 encoded")); - } - let payload = payload.trim(); - if payload.is_empty() { - return Err(Error::invalid_params().data("image data is empty")); - } - - Ok((mime_type, payload)) -} - #[cfg(test)] mod tests { use super::super::test_support::test_session; use super::*; use agent_client_protocol::schema::v1::ContentBlock; use agent_client_protocol::schema::v1::EmbeddedResourceResource; - use std::path::Path; #[test] fn flattens_text_and_embedded_context() { @@ -259,17 +145,6 @@ mod tests { assert_eq!(text, "Inspect this.[file:///tmp/main.rs]\nfn main() {}"); } - #[test] - fn writes_supported_acp_image_to_managed_session_storage() { - let session_id = format!("acp-image-{}", cuid2::create_id()); - let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/png"); - let path = write_prompt_image(&session_id, &image).expect("image file"); - - assert!(Path::new(&path).starts_with(crate::persistence::attachments::root_dir())); - assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); - crate::persistence::attachments::cleanup_session(&session_id).unwrap(); - } - #[test] fn command_text_override_preserves_embedded_context() { let resource = agent_client_protocol::schema::v1::EmbeddedResource::new( @@ -306,17 +181,6 @@ mod tests { assert!(audio.is_empty()); } - #[test] - fn writes_supported_acp_audio_to_managed_session_storage() { - let session_id = format!("acp-audio-{}", cuid2::create_id()); - let audio = agent_client_protocol::schema::v1::AudioContent::new("YXVkaW8=", "audio/wav"); - let path = write_prompt_audio(&session_id, &audio).expect("audio file"); - - assert!(path.ends_with(".wav")); - assert_eq!(std::fs::read(&path).unwrap(), b"audio"); - crate::persistence::attachments::cleanup_session(&session_id).unwrap(); - } - #[test] fn rejects_audio_for_models_without_audio_modality() { let session_id = format!("acp-audio-{}", cuid2::create_id()); @@ -335,20 +199,6 @@ mod tests { crate::persistence::attachments::cleanup_session(&session_id).unwrap(); } - #[test] - fn writes_acp_clipboard_image_data_uri_to_managed_storage() { - let session_id = format!("acp-image-{}", cuid2::create_id()); - let image = agent_client_protocol::schema::v1::ImageContent::new( - "data:image/png;base64,aGk=", - "application/octet-stream", - ); - let path = write_prompt_image(&session_id, &image).expect("image file"); - - assert!(path.ends_with(".png")); - assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); - crate::persistence::attachments::cleanup_session(&session_id).unwrap(); - } - #[test] fn prompt_image_failure_rolls_back_prior_managed_files() { let session_id = format!("acp-image-{}", cuid2::create_id()); @@ -380,21 +230,4 @@ mod tests { ); crate::persistence::attachments::cleanup_session(&session_id).unwrap(); } - - #[test] - fn rejects_non_base64_acp_image_data_uri() { - let image = agent_client_protocol::schema::v1::ImageContent::new( - "data:image/png,not-base64", - "image/png", - ); - - assert!(write_prompt_image("test", &image).is_err()); - } - - #[test] - fn rejects_unsupported_acp_image_mime_type() { - let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/tiff"); - - assert!(write_prompt_image("test", &image).is_err()); - } } diff --git a/src/acp/service/command_expansion.rs b/src/acp/service/command_expansion.rs new file mode 100644 index 0000000..401ec17 --- /dev/null +++ b/src/acp/service/command_expansion.rs @@ -0,0 +1,236 @@ +use super::internal_error; +use super::AcpSession; +use agent_client_protocol::Error; + +#[derive(Debug, PartialEq, Eq)] +pub(super) enum SlashExpansion { + Prompt { + prompt: String, + agent: Option, + model: Option, + }, + LocalResult(String), + Btw(String), +} + +pub(super) async fn expand_slash_command( + session: &AcpSession, + prompt: &str, +) -> Result { + let Some(command_line) = prompt.strip_prefix('/') else { + return Ok(SlashExpansion::Prompt { + prompt: prompt.to_string(), + agent: None, + model: None, + }); + }; + let (name, args) = command_line + .split_once(char::is_whitespace) + .map(|(name, args)| (name, args.trim_start())) + .unwrap_or((command_line, "")); + if name == "btw" { + let question = args.trim(); + if question.is_empty() { + return Err(Error::invalid_params().data("Usage: /btw ")); + } + return Ok(SlashExpansion::Btw(question.to_string())); + } + if name == "skills" { + if !args.is_empty() { + return Err(Error::invalid_params().data("Usage: /skills")); + } + let skills = session + .skills + .all() + .into_iter() + .map(|skill| { + format!( + "- /{} — {}", + skill.name, + skill.description.as_deref().unwrap_or("No description") + ) + }) + .collect::>(); + return Ok(SlashExpansion::LocalResult(if skills.is_empty() { + "No skills are available in this workspace.".to_string() + } else { + format!("Available skills:\n{}", skills.join("\n")) + })); + } + if name == "mcp" { + if !args.is_empty() { + return Err(Error::invalid_params().data("Usage: /mcp")); + } + let manager = crate::mcp::McpManager::ensure( + session.config.merged_config.mcp.clone(), + session.cwd.clone(), + ); + let servers = manager + .lock() + .await + .views() + .into_iter() + .map(|server| { + let detail = server + .detail + .map(|detail| format!(": {detail}")) + .unwrap_or_default(); + format!( + "- {} ({}, {}){}", + server.name, server.kind, server.status, detail + ) + }) + .collect::>(); + return Ok(SlashExpansion::LocalResult(if servers.is_empty() { + "No MCP servers are configured for this workspace.".to_string() + } else { + format!("MCP servers:\n{}", servers.join("\n")) + })); + } + if let Some(command) = session + .config + .merged_config + .commands + .iter() + .find(|command| command.name == name) + { + return command + .render(args) + .await + .map_err(|_| internal_error()) + .map(|rendered| SlashExpansion::Prompt { + prompt: rendered.prompt, + agent: rendered.agent, + model: rendered.model, + }); + } + if let Some(skill) = session.skills.get(name) { + let mut expanded = skill.content.clone(); + if !args.is_empty() { + expanded.push_str("\n\nUser task/context:\n"); + expanded.push_str(args); + } + return Ok(SlashExpansion::Prompt { + prompt: expanded, + agent: None, + model: None, + }); + } + Err(Error::invalid_params().data(format!("Unknown ACP command: /{name}"))) +} + +#[cfg(test)] +mod tests { + use super::super::commands::available_commands; + use super::super::test_support::config_with_command; + + use super::super::test_support::session_with_config; + use super::*; + use std::path::PathBuf; + + #[tokio::test] + async fn expands_custom_slash_command_before_prompting() { + let config = config_with_command(crate::command::custom::CustomCommand { + name: "review".to_string(), + description: None, + template: "Review this carefully: $ARGUMENTS".to_string(), + agent: Some("build".to_string()), + model: Some("openai/gpt-5".to_string()), + subtask: Some(false), + source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( + "/tmp/opencode.jsonc", + )), + workdir: PathBuf::from("/tmp"), + }); + + let session = session_with_config(config); + let prompt = expand_slash_command(&session, "/review src/acp/service.rs") + .await + .expect("expanded command"); + + assert_eq!( + prompt, + SlashExpansion::Prompt { + prompt: "Review this carefully: src/acp/service.rs".to_string(), + agent: Some("build".to_string()), + model: Some("openai/gpt-5".to_string()), + } + ); + } + + #[tokio::test] + async fn advertises_and_expands_workspace_skills() { + let temp = tempfile::tempdir().expect("temp dir"); + let skill_dir = temp.path().join(".opencode/skill/reviewer"); + std::fs::create_dir_all(&skill_dir).expect("skill dir"); + std::fs::write( + skill_dir.join("SKILL.md"), + "---\nname: reviewer\ndescription: Review code carefully\n---\nInspect correctness and risks.", + ) + .expect("skill file"); + let mut config = config_with_command(crate::command::custom::CustomCommand { + name: "custom".to_string(), + description: None, + template: "$ARGUMENTS".to_string(), + agent: None, + model: None, + subtask: Some(false), + source: crate::command::custom::CustomCommandSource::Config( + temp.path().join("opencode.jsonc"), + ), + workdir: temp.path().to_path_buf(), + }); + config.project_root = temp.path().to_path_buf(); + config.cwd = temp.path().to_path_buf(); + config.xdg_config_home = temp.path().join("config"); + let session = session_with_config(config); + + let commands = available_commands(&session); + assert!(commands.iter().any(|command| command.name == "reviewer")); + let prompt = expand_slash_command(&session, "/reviewer src/lib.rs") + .await + .expect("expanded skill"); + let SlashExpansion::Prompt { prompt, .. } = prompt else { + panic!("expected skill prompt"); + }; + assert!(prompt.contains("Inspect correctness and risks.")); + assert!(prompt.contains("src/lib.rs")); + } + + #[tokio::test] + async fn reports_configured_mcp_servers() { + let mut config = config_with_command(crate::command::custom::CustomCommand { + name: "custom".to_string(), + description: None, + template: "$ARGUMENTS".to_string(), + agent: None, + model: None, + subtask: Some(false), + source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( + "/tmp/opencode.jsonc", + )), + workdir: PathBuf::from("/tmp"), + }); + config.merged_config.mcp.insert( + "filesystem".to_string(), + crate::config::configuration::McpServerConfig::Local( + crate::config::configuration::McpLocalConfig { + command: vec!["mcp-server".to_string()], + cwd: None, + environment: Default::default(), + enabled: true, + timeout_ms: None, + }, + ), + ); + let session = session_with_config(config); + + let prompt = expand_slash_command(&session, "/mcp") + .await + .expect("mcp status"); + let SlashExpansion::LocalResult(prompt) = prompt else { + panic!("expected local MCP result"); + }; + assert!(prompt.contains("filesystem (local, connecting)")); + } +} diff --git a/src/acp/service/commands.rs b/src/acp/service/commands.rs index 30713f8..f12fbec 100644 --- a/src/acp/service/commands.rs +++ b/src/acp/service/commands.rs @@ -44,7 +44,7 @@ pub(super) fn compact_command(prompt: &str) -> Result { Ok(true) } -fn available_commands(session: &AcpSession) -> Vec { +pub(super) fn available_commands(session: &AcpSession) -> Vec { let mut commands: Vec<_> = session .config .merged_config @@ -106,123 +106,6 @@ fn available_commands(session: &AcpSession) -> Vec { commands } -#[derive(Debug, PartialEq, Eq)] -pub(super) enum SlashExpansion { - Prompt { - prompt: String, - agent: Option, - model: Option, - }, - LocalResult(String), - Btw(String), -} - -pub(super) async fn expand_slash_command( - session: &AcpSession, - prompt: &str, -) -> Result { - let Some(command_line) = prompt.strip_prefix('/') else { - return Ok(SlashExpansion::Prompt { - prompt: prompt.to_string(), - agent: None, - model: None, - }); - }; - let (name, args) = command_line - .split_once(char::is_whitespace) - .map(|(name, args)| (name, args.trim_start())) - .unwrap_or((command_line, "")); - if name == "btw" { - let question = args.trim(); - if question.is_empty() { - return Err(Error::invalid_params().data("Usage: /btw ")); - } - return Ok(SlashExpansion::Btw(question.to_string())); - } - if name == "skills" { - if !args.is_empty() { - return Err(Error::invalid_params().data("Usage: /skills")); - } - let skills = session - .skills - .all() - .into_iter() - .map(|skill| { - format!( - "- /{} — {}", - skill.name, - skill.description.as_deref().unwrap_or("No description") - ) - }) - .collect::>(); - return Ok(SlashExpansion::LocalResult(if skills.is_empty() { - "No skills are available in this workspace.".to_string() - } else { - format!("Available skills:\n{}", skills.join("\n")) - })); - } - if name == "mcp" { - if !args.is_empty() { - return Err(Error::invalid_params().data("Usage: /mcp")); - } - let manager = crate::mcp::McpManager::ensure( - session.config.merged_config.mcp.clone(), - session.cwd.clone(), - ); - let servers = manager - .lock() - .await - .views() - .into_iter() - .map(|server| { - let detail = server - .detail - .map(|detail| format!(": {detail}")) - .unwrap_or_default(); - format!( - "- {} ({}, {}){}", - server.name, server.kind, server.status, detail - ) - }) - .collect::>(); - return Ok(SlashExpansion::LocalResult(if servers.is_empty() { - "No MCP servers are configured for this workspace.".to_string() - } else { - format!("MCP servers:\n{}", servers.join("\n")) - })); - } - if let Some(command) = session - .config - .merged_config - .commands - .iter() - .find(|command| command.name == name) - { - return command - .render(args) - .await - .map_err(|_| internal_error()) - .map(|rendered| SlashExpansion::Prompt { - prompt: rendered.prompt, - agent: rendered.agent, - model: rendered.model, - }); - } - if let Some(skill) = session.skills.get(name) { - let mut expanded = skill.content.clone(); - if !args.is_empty() { - expanded.push_str("\n\nUser task/context:\n"); - expanded.push_str(args); - } - return Ok(SlashExpansion::Prompt { - prompt: expanded, - agent: None, - model: None, - }); - } - Err(Error::invalid_params().data(format!("Unknown ACP command: /{name}"))) -} - impl AcpService { pub async fn available_commands( &self, @@ -279,11 +162,15 @@ impl AcpService { #[cfg(test)] mod tests { + use super::super::command_expansion::expand_slash_command; + use super::super::command_expansion::SlashExpansion; use super::super::test_support::config_with_command; + use super::super::test_support::session_with_config; use super::super::test_support::test_session; use super::*; use agent_client_protocol::schema::v1::AvailableCommandInput; + use std::path::PathBuf; #[test] @@ -356,110 +243,4 @@ mod tests { Some(AvailableCommandInput::Unstructured(_)) )); } - - #[tokio::test] - async fn expands_custom_slash_command_before_prompting() { - let config = config_with_command(crate::command::custom::CustomCommand { - name: "review".to_string(), - description: None, - template: "Review this carefully: $ARGUMENTS".to_string(), - agent: Some("build".to_string()), - model: Some("openai/gpt-5".to_string()), - subtask: Some(false), - source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( - "/tmp/opencode.jsonc", - )), - workdir: PathBuf::from("/tmp"), - }); - - let session = session_with_config(config); - let prompt = expand_slash_command(&session, "/review src/acp/service.rs") - .await - .expect("expanded command"); - - assert_eq!( - prompt, - SlashExpansion::Prompt { - prompt: "Review this carefully: src/acp/service.rs".to_string(), - agent: Some("build".to_string()), - model: Some("openai/gpt-5".to_string()), - } - ); - } - - #[tokio::test] - async fn advertises_and_expands_workspace_skills() { - let temp = tempfile::tempdir().expect("temp dir"); - let skill_dir = temp.path().join(".opencode/skill/reviewer"); - std::fs::create_dir_all(&skill_dir).expect("skill dir"); - std::fs::write( - skill_dir.join("SKILL.md"), - "---\nname: reviewer\ndescription: Review code carefully\n---\nInspect correctness and risks.", - ) - .expect("skill file"); - let mut config = config_with_command(crate::command::custom::CustomCommand { - name: "custom".to_string(), - description: None, - template: "$ARGUMENTS".to_string(), - agent: None, - model: None, - subtask: Some(false), - source: crate::command::custom::CustomCommandSource::Config( - temp.path().join("opencode.jsonc"), - ), - workdir: temp.path().to_path_buf(), - }); - config.project_root = temp.path().to_path_buf(); - config.cwd = temp.path().to_path_buf(); - config.xdg_config_home = temp.path().join("config"); - let session = session_with_config(config); - - let commands = available_commands(&session); - assert!(commands.iter().any(|command| command.name == "reviewer")); - let prompt = expand_slash_command(&session, "/reviewer src/lib.rs") - .await - .expect("expanded skill"); - let SlashExpansion::Prompt { prompt, .. } = prompt else { - panic!("expected skill prompt"); - }; - assert!(prompt.contains("Inspect correctness and risks.")); - assert!(prompt.contains("src/lib.rs")); - } - - #[tokio::test] - async fn reports_configured_mcp_servers() { - let mut config = config_with_command(crate::command::custom::CustomCommand { - name: "custom".to_string(), - description: None, - template: "$ARGUMENTS".to_string(), - agent: None, - model: None, - subtask: Some(false), - source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( - "/tmp/opencode.jsonc", - )), - workdir: PathBuf::from("/tmp"), - }); - config.merged_config.mcp.insert( - "filesystem".to_string(), - crate::config::configuration::McpServerConfig::Local( - crate::config::configuration::McpLocalConfig { - command: vec!["mcp-server".to_string()], - cwd: None, - environment: Default::default(), - enabled: true, - timeout_ms: None, - }, - ), - ); - let session = session_with_config(config); - - let prompt = expand_slash_command(&session, "/mcp") - .await - .expect("mcp status"); - let SlashExpansion::LocalResult(prompt) = prompt else { - panic!("expected local MCP result"); - }; - assert!(prompt.contains("filesystem (local, connecting)")); - } } diff --git a/src/acp/service/compaction.rs b/src/acp/service/compaction.rs index 10b1cf0..7a8e2ff 100644 --- a/src/acp/service/compaction.rs +++ b/src/acp/service/compaction.rs @@ -1,12 +1,12 @@ use super::cancellable; -use super::config::model_output_limit; -use super::config::model_reasoning_capability; +use super::compaction_summary::compacted_messages; +use super::compaction_summary::compaction_reasoning; +use super::context::acp_context_tokens; use super::internal_error; use super::internal_error_with; +use super::models::model_output_limit; use super::notifications::send_text; -use super::usage::acp_context_tokens; use super::usage::acp_session_cost; -use super::usage::estimate_session_usage_cost; use super::usage::send_usage; use super::AcpService; use super::AcpSession; @@ -17,87 +17,6 @@ use agent_client_protocol::ConnectionTo; use agent_client_protocol::Error; use tokio_util::sync::CancellationToken; -fn compaction_reasoning(session: &AcpSession) -> Option { - use crate::model::reasoning::ReasoningEffort; - let capability = model_reasoning_capability( - &session.config, - &session.models, - &session.provider, - &session.model, - )?; - [ - ReasoningEffort::None, - ReasoningEffort::Minimal, - ReasoningEffort::Low, - ] - .into_iter() - .find(|effort| capability.values().contains(effort)) - .or(session.reasoning) - .filter(|effort| *effort != ReasoningEffort::None) -} - -fn compacted_messages( - messages: &[crate::session::types::Message], - selection: &crate::session::compaction::CompactionSelection, - summary: &crate::llm::client::CompactionSummary, - session: &AcpSession, - before_tokens: usize, - before_messages: usize, -) -> Result< - ( - Vec, - crate::session::types::CompactionStats, - ), - Error, -> { - let mut compacted = crate::session::compaction::apply_soft_compaction( - messages, - selection, - &summary.text, - Some(session.model.clone()), - Some(session.provider.clone()), - Some(session.agent.clone()), - crate::session::types::CompactionStats { - before_tokens, - after_tokens: 0, - before_messages, - after_messages: 0, - }, - ); - crate::session::compaction::attach_summary_usage( - &mut compacted, - crate::session::types::RecordedUsage { - input: summary.usage.input, - output: summary.usage.output, - cache_read: summary.usage.cache_read, - cache_write: summary.usage.cache_write, - cost: estimate_session_usage_cost(session, &summary.usage), - }, - ); - let after_tokens = crate::session::compaction::total_context_tokens(&compacted); - let after_messages = crate::session::compaction::filter_messages_for_context(&compacted).len(); - let stats = crate::session::types::CompactionStats { - before_tokens, - after_tokens, - before_messages, - after_messages, - }; - if after_tokens >= before_tokens { - return Err(Error::invalid_params().data(format!( - "Compaction did not reduce context ({})", - crate::session::compaction::format_compaction_stats(stats) - ))); - } - if let Some(marker) = compacted - .iter_mut() - .rev() - .find(|message| crate::session::compaction::is_compaction_marker(message)) - { - marker.compaction_stats = Some(stats); - } - Ok((compacted, stats)) -} - impl AcpService { pub(super) async fn compact_session( &self, @@ -319,292 +238,5 @@ impl AcpService { } #[cfg(test)] -mod tests { - use super::super::prompt::acp_request_messages; - use super::super::test_support::test_session; - use super::super::usage::acp_context_tokens; - use super::super::usage::acp_session_cost; - use super::super::AcpService; - use super::super::AcpSessions; - use super::*; - use crate::session::manager::SessionManager; - use agent_client_protocol::schema::v1::SessionNotification; - use agent_client_protocol::schema::v1::SessionUpdate; - use agent_client_protocol::Client; - use std::sync::Arc; - use std::sync::Mutex; - use tokio::sync::Mutex as AsyncMutex; - - #[test] - fn acp_compaction_discards_provider_usage_before_compaction_boundary() { - use crate::session::types::{CompactionStats, Message, MessagePart}; - let mut assistant = Message::assistant("old answer"); - assistant - .parts - .push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); - let mut summary = Message::user(format!( - "{}\nsummary", - crate::session::compaction::SUMMARY_PREFIX - )); - summary - .parts - .push(MessagePart::usage(95_000, 500, 0, 0, 0.0)); - let marker = crate::session::compaction::compaction_marker(CompactionStats { - before_tokens: 90_100, - after_tokens: 100, - before_messages: 2, - after_messages: 1, - }); - let messages = vec![Message::user("question"), assistant, summary, marker]; - - assert_eq!( - acp_context_tokens(&messages, 1_000), - crate::session::compaction::total_context_tokens(&messages) + 1_000 - ); - } - - #[test] - fn acp_compaction_honors_reserved_tokens_and_auto_setting() { - use crate::config::configuration::CompactionConfig; - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant - .parts - .push(MessagePart::usage(70_000, 0, 10_000, 0, 0.0)); - let messages = vec![Message::user("history ".repeat(2_000)), assistant]; - assert!(crate::session::compaction::select_messages_for_compaction( - &messages, - crate::session::compaction::DEFAULT_TAIL_TURNS, - ) - .is_some()); - let config = CompactionConfig::Settings { - auto: true, - prune: false, - reserved: Some(20_000), - }; - let used = acp_context_tokens(&messages, 1_000); - assert!(crate::session::compaction::total_context_tokens(&messages) < 80_000); - assert!(crate::session::compaction::should_auto_compact( - &config, - used, - Some(100_000), - None - )); - assert!(!crate::session::compaction::should_auto_compact( - &config, - used, - Some(100_001), - None - )); - assert!(!crate::session::compaction::should_auto_compact( - &CompactionConfig::Disabled, - used, - Some(100_000), - None - )); - assert!(!crate::session::compaction::should_auto_compact( - &config, used, None, None - )); - let no_auto = CompactionConfig::Settings { - auto: false, - prune: false, - reserved: Some(20_000), - }; - assert!(!crate::session::compaction::should_auto_compact( - &no_auto, - used, - Some(100_000), - None - )); - } - - #[test] - fn acp_compaction_keeps_fresh_system_prompt_in_request() { - use crate::session::types::{CompactionStats, Message}; - let summary = Message::user(format!( - "{}\nsummary", - crate::session::compaction::SUMMARY_PREFIX - )); - let tail = Message::assistant("retained answer"); - let marker = crate::session::compaction::compaction_marker(CompactionStats { - before_tokens: 10_000, - after_tokens: 100, - before_messages: 3, - after_messages: 2, - }); - let pending = Message::user("continue"); - let history = vec![ - Message::user("archived history"), - summary.clone(), - tail.clone(), - marker, - pending.clone(), - ]; - let system = Message::system("fresh system instructions"); - let request = acp_request_messages(&history, system.clone()); - assert_eq!(request, vec![system, summary, tail, pending]); - assert_eq!( - crate::session::compaction::filter_messages_for_context(&request), - request - ); - } - - #[test] - fn acp_compaction_cost_keeps_summary_billing_without_double_counting() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant.cost = Some(0.25); - assistant - .parts - .push(MessagePart::usage(100, 20, 0, 0, 0.25)); - let mut summary = Message::user(format!( - "{}\nsummary", - crate::session::compaction::SUMMARY_PREFIX - )); - summary.parts.push(MessagePart::usage(100, 20, 0, 0, 0.05)); - assert!((acp_session_cost(&[assistant, summary]) - 0.30).abs() < f64::EPSILON); - } - - #[tokio::test(flavor = "current_thread")] - async fn acp_compaction_notifies_client_of_reduced_context_and_summary_cost() { - use crate::session::types::{Message, MessagePart}; - use agent_client_protocol::Agent; - tokio::task::LocalSet::new() - .run_until(async { - let workspace = tempfile::tempdir().unwrap(); - let mut manager = SessionManager::new() - .with_history_for_workspace(workspace.path()) - .unwrap(); - let session_id = manager.create_session(Some("ACP compaction regression".into())); - let mut assistant = Message::assistant("old answer"); - assistant - .parts - .push(MessagePart::usage(80_000, 200, 0, 0, 0.25)); - let messages = vec![Message::user("history ".repeat(2_000)), assistant]; - let selection = crate::session::compaction::select_messages_for_compaction( - &messages, - crate::session::compaction::DEFAULT_TAIL_TURNS, - ) - .unwrap(); - let mut session = test_session(); - session.context_window = Some(100_000); - let summary = crate::llm::client::CompactionSummary { - text: "Retained task details".into(), - usage: crate::aisdk::chunk::TokenUsage { - input: 3_000, - output: 20, - cache_read: 0, - cache_write: 0, - }, - }; - let (mut compacted, stats) = compacted_messages( - &messages, - &selection, - &summary, - &session, - crate::session::compaction::total_context_tokens(&messages), - messages.len(), - ) - .unwrap(); - let summary_message = compacted - .iter_mut() - .find(|message| crate::session::compaction::is_compaction_summary(message)) - .unwrap(); - summary_message - .parts - .push(MessagePart::usage(0, 1, 0, 0, 0.05)); - let expected_used = - crate::session::compaction::total_context_tokens(&compacted) + 1_000; - manager - .replace_session_messages(&session_id, compacted) - .unwrap(); - let service = AcpService { - sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), - session_manager: Arc::new(Mutex::new(manager)), - client_capabilities: Arc::new(Mutex::new(Default::default())), - }; - let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); - let client = Client.builder().on_receive_notification( - async move |notification: SessionNotification, _connection| { - let _ = observed_tx.send(notification.update); - Ok(()) - }, - agent_client_protocol::on_receive_notification!(), - ); - Agent - .builder() - .connect_with(client, async move |connection| { - service - .notify_compaction(&session_id, &session, &connection, stats, 1_000) - .unwrap(); - let feedback = tokio::time::timeout( - std::time::Duration::from_secs(2), - observed_rx.recv(), - ) - .await - .unwrap() - .unwrap(); - assert!(matches!(feedback, SessionUpdate::AgentMessageChunk(_))); - let update = tokio::time::timeout( - std::time::Duration::from_secs(2), - observed_rx.recv(), - ) - .await - .unwrap() - .unwrap(); - let SessionUpdate::UsageUpdate(usage) = update else { - panic!("missing usage update") - }; - assert_eq!(usage.used, expected_used as u64); - assert_eq!(usage.size, 100_000); - assert!((usage.cost.unwrap().amount - 0.30).abs() < f64::EPSILON); - service - .session_manager - .lock() - .unwrap() - .try_delete_session(&session_id) - .unwrap(); - Ok(()) - }) - .await - .unwrap(); - }) - .await; - } - - #[test] - fn builds_smaller_soft_compaction_for_acp() { - let session = test_session(); - let messages = vec![ - crate::session::types::Message::user("u".repeat(8_000)), - crate::session::types::Message::assistant("a".repeat(8_000)), - crate::session::types::Message::user("recent"), - ]; - let selection = crate::session::compaction::select_messages_for_compaction_with_min( - &messages, - crate::session::compaction::DEFAULT_TAIL_TURNS, - 0, - ) - .expect("compaction selection"); - let before_tokens = crate::session::compaction::total_context_tokens(&messages); - let before_messages = - crate::session::compaction::filter_messages_for_context(&messages).len(); - - let (compacted, stats) = compacted_messages( - &messages, - &selection, - &crate::llm::client::CompactionSummary { - text: "short handoff".to_string(), - usage: Default::default(), - }, - &session, - before_tokens, - before_messages, - ) - .expect("smaller compaction"); - - assert!(stats.after_tokens < stats.before_tokens); - assert!(crate::session::compaction::latest_compaction_stats(&compacted).is_some()); - assert_eq!(compacted[0].id, messages[0].id); - } -} +#[path = "tests/compaction.rs"] +mod tests; diff --git a/src/acp/service/compaction_summary.rs b/src/acp/service/compaction_summary.rs new file mode 100644 index 0000000..ad4df6d --- /dev/null +++ b/src/acp/service/compaction_summary.rs @@ -0,0 +1,129 @@ +use super::models::model_reasoning_capability; +use super::usage::estimate_session_usage_cost; +use super::AcpSession; +use agent_client_protocol::Error; + +pub(super) fn compaction_reasoning( + session: &AcpSession, +) -> Option { + use crate::model::reasoning::ReasoningEffort; + let capability = model_reasoning_capability( + &session.config, + &session.models, + &session.provider, + &session.model, + )?; + [ + ReasoningEffort::None, + ReasoningEffort::Minimal, + ReasoningEffort::Low, + ] + .into_iter() + .find(|effort| capability.values().contains(effort)) + .or(session.reasoning) + .filter(|effort| *effort != ReasoningEffort::None) +} + +pub(super) fn compacted_messages( + messages: &[crate::session::types::Message], + selection: &crate::session::compaction::CompactionSelection, + summary: &crate::llm::client::CompactionSummary, + session: &AcpSession, + before_tokens: usize, + before_messages: usize, +) -> Result< + ( + Vec, + crate::session::types::CompactionStats, + ), + Error, +> { + let mut compacted = crate::session::compaction::apply_soft_compaction( + messages, + selection, + &summary.text, + Some(session.model.clone()), + Some(session.provider.clone()), + Some(session.agent.clone()), + crate::session::types::CompactionStats { + before_tokens, + after_tokens: 0, + before_messages, + after_messages: 0, + }, + ); + crate::session::compaction::attach_summary_usage( + &mut compacted, + crate::session::types::RecordedUsage { + input: summary.usage.input, + output: summary.usage.output, + cache_read: summary.usage.cache_read, + cache_write: summary.usage.cache_write, + cost: estimate_session_usage_cost(session, &summary.usage), + }, + ); + let after_tokens = crate::session::compaction::total_context_tokens(&compacted); + let after_messages = crate::session::compaction::filter_messages_for_context(&compacted).len(); + let stats = crate::session::types::CompactionStats { + before_tokens, + after_tokens, + before_messages, + after_messages, + }; + if after_tokens >= before_tokens { + return Err(Error::invalid_params().data(format!( + "Compaction did not reduce context ({})", + crate::session::compaction::format_compaction_stats(stats) + ))); + } + if let Some(marker) = compacted + .iter_mut() + .rev() + .find(|message| crate::session::compaction::is_compaction_marker(message)) + { + marker.compaction_stats = Some(stats); + } + Ok((compacted, stats)) +} + +#[cfg(test)] +mod tests { + use super::super::test_support::test_session; + use super::*; + + #[test] + fn builds_smaller_soft_compaction_for_acp() { + let session = test_session(); + let messages = vec![ + crate::session::types::Message::user("u".repeat(8_000)), + crate::session::types::Message::assistant("a".repeat(8_000)), + crate::session::types::Message::user("recent"), + ]; + let selection = crate::session::compaction::select_messages_for_compaction_with_min( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + 0, + ) + .expect("compaction selection"); + let before_tokens = crate::session::compaction::total_context_tokens(&messages); + let before_messages = + crate::session::compaction::filter_messages_for_context(&messages).len(); + + let (compacted, stats) = compacted_messages( + &messages, + &selection, + &crate::llm::client::CompactionSummary { + text: "short handoff".to_string(), + usage: Default::default(), + }, + &session, + before_tokens, + before_messages, + ) + .expect("smaller compaction"); + + assert!(stats.after_tokens < stats.before_tokens); + assert!(crate::session::compaction::latest_compaction_stats(&compacted).is_some()); + assert_eq!(compacted[0].id, messages[0].id); + } +} diff --git a/src/acp/service/config.rs b/src/acp/service/config.rs index 0cd976e..1cc0ab9 100644 --- a/src/acp/service/config.rs +++ b/src/acp/service/config.rs @@ -1,170 +1,17 @@ +use super::models::find_selectable_model; +use super::models::model_config_option; +use super::models::model_context_window; +use super::models::reasoning_config_option; +use super::models::resolved_reasoning; +use super::modes::mode_description; use super::AcpService; use super::AcpSession; -use crate::config::configuration::LoadedConfig; -use agent_client_protocol::schema::v1::McpServer; use agent_client_protocol::schema::v1::SessionConfigOption; use agent_client_protocol::schema::v1::SessionConfigOptionCategory; -use agent_client_protocol::schema::v1::SessionConfigSelectGroup; use agent_client_protocol::schema::v1::SessionConfigSelectOption; -use agent_client_protocol::schema::v1::SessionMode; -use agent_client_protocol::schema::v1::SessionModeState; use agent_client_protocol::schema::v1::SetSessionConfigOptionResponse; use agent_client_protocol::Error; -pub(super) fn model_output_limit( - config: &LoadedConfig, - provider_id: &str, - model_id: &str, -) -> Option { - if let Some(limit) = config - .merged_config - .custom_providers - .get(&provider_id.trim().to_ascii_lowercase()) - .and_then(|provider| provider.models.get(model_id)) - .and_then(|model| model.max_tokens) - { - return Some(limit); - } - crate::model::discovery::Discovery::new_with_config( - Some(config.merged_config.custom_providers.clone()), - config.merged_config.disabled_providers.clone(), - config.merged_config.enabled_providers.clone(), - ) - .ok() - .and_then(|discovery| discovery.get_model_output_limit(provider_id, model_id)) -} - -pub(super) fn resolved_reasoning( - session: &AcpSession, - requested: crate::model::reasoning::ReasoningEffort, -) -> Option { - model_reasoning_capability( - &session.config, - &session.models, - &session.provider, - &session.model, - ) - .and_then(|capability| capability.resolve(Some(requested))) -} - -fn reasoning_effort_label(effort: crate::model::reasoning::ReasoningEffort) -> &'static str { - match effort { - crate::model::reasoning::ReasoningEffort::None => "None", - crate::model::reasoning::ReasoningEffort::Minimal => "Minimal", - crate::model::reasoning::ReasoningEffort::Low => "Low", - crate::model::reasoning::ReasoningEffort::Medium => "Medium", - crate::model::reasoning::ReasoningEffort::High => "High", - crate::model::reasoning::ReasoningEffort::XHigh => "Xhigh", - crate::model::reasoning::ReasoningEffort::Max => "Max", - } -} - -pub(super) fn merge_acp_mcp_servers(config: &mut LoadedConfig, servers: Vec) { - for server in servers { - let (name, server) = match server { - McpServer::Stdio(server) => { - let mut command = vec![server.command.to_string_lossy().into_owned()]; - command.extend(server.args); - ( - server.name, - crate::config::configuration::McpServerConfig::Local( - crate::config::configuration::McpLocalConfig { - command, - cwd: None, - environment: server - .env - .into_iter() - .map(|variable| (variable.name, variable.value)) - .collect(), - enabled: true, - timeout_ms: None, - }, - ), - ) - } - McpServer::Http(server) => ( - server.name, - crate::config::configuration::McpServerConfig::Remote( - crate::config::configuration::McpRemoteConfig { - url: server.url, - headers: server - .headers - .into_iter() - .map(|header| (header.name, header.value)) - .collect(), - enabled: true, - timeout_ms: None, - oauth_enabled: false, - oauth_client_id: None, - oauth_client_secret: None, - oauth_scope: None, - }, - ), - ), - McpServer::Sse(server) => ( - server.name, - crate::config::configuration::McpServerConfig::Remote( - crate::config::configuration::McpRemoteConfig { - url: server.url, - headers: server - .headers - .into_iter() - .map(|header| (header.name, header.value)) - .collect(), - enabled: true, - timeout_ms: None, - oauth_enabled: false, - oauth_client_id: None, - oauth_client_secret: None, - oauth_scope: None, - }, - ), - ), - #[allow(unreachable_patterns)] - _ => continue, - }; - config.merged_config.mcp.insert(name, server); - } -} - -pub(super) fn resolve_model(config: &LoadedConfig) -> (String, String) { - if let Ok(Some((provider, model))) = - crate::persistence::PrefsDAO::new().and_then(|prefs| prefs.get_active_model()) - { - return (provider, model); - } - config - .merged_config - .model - .as_deref() - .map(crate::app::parse_model_ref) - .unwrap_or_else(|| ("opencode".to_string(), "big-pickle".to_string())) -} - -fn mode_description(agent: &crate::agent::definition::AgentDefinition) -> &str { - if agent.name == "compaction" && agent.description.trim().is_empty() { - "Compaction agent. Summarizes conversation context for continued work." - } else { - &agent.description - } -} - -pub(super) fn session_modes(session: &AcpSession) -> SessionModeState { - let mut modes = session - .config - .merged_config - .agent_registry - .visible_primary_agents() - .into_iter() - .map(|agent| { - SessionMode::new(agent.name.clone(), agent.name.clone()) - .description(mode_description(agent)) - }) - .collect::>(); - modes.sort_by(|left, right| left.name.cmp(&right.name)); - SessionModeState::new(session.agent.clone(), modes) -} - pub(super) fn session_config_options(session: &AcpSession) -> Vec { let mut mode_options = session .config @@ -188,137 +35,6 @@ pub(super) fn session_config_options(session: &AcpSession) -> Vec SessionConfigOption { - let options = [ - crate::model::reasoning::ReasoningEffort::None, - crate::model::reasoning::ReasoningEffort::Low, - crate::model::reasoning::ReasoningEffort::Medium, - crate::model::reasoning::ReasoningEffort::High, - crate::model::reasoning::ReasoningEffort::XHigh, - crate::model::reasoning::ReasoningEffort::Max, - ] - .iter() - .map(|effort| SessionConfigSelectOption::new(effort.as_str(), reasoning_effort_label(*effort))) - .collect::>(); - let current = session.reasoning_selection.as_str(); - SessionConfigOption::select("effort", "Effort", current, options) - .category(SessionConfigOptionCategory::ThoughtLevel) -} - -fn model_config_option( - models: &[crate::model::types::Model], - provider: &str, - model_id: &str, -) -> SessionConfigOption { - let mut model_groups: Vec = Vec::new(); - for model in models { - let option = SessionConfigSelectOption::new(model_value(model), model.name.clone()) - .description(model.id.clone()); - if model_groups - .last() - .is_some_and(|group| group.group.to_string() == model.provider_id) - { - model_groups - .last_mut() - .expect("model group exists") - .options - .push(option); - } else { - model_groups.push(SessionConfigSelectGroup::new( - model.provider_id.clone(), - model.provider_name.clone(), - vec![option], - )); - } - } - - SessionConfigOption::select( - "model", - "Model", - format!("{provider}/{model_id}"), - model_groups, - ) - .category(SessionConfigOptionCategory::Model) -} - -fn model_value(model: &crate::model::types::Model) -> String { - crate::model::catalog::model_ref(model) -} - -pub(super) fn find_selectable_model<'a>( - models: &'a [crate::model::types::Model], - model_ref: &str, -) -> Result<&'a crate::model::types::Model, Error> { - models - .iter() - .find(|model| model_value(model) == model_ref) - .ok_or_else(|| Error::invalid_params().data("unknown or unavailable ACP model")) -} - -pub(super) fn model_reasoning( - config: &LoadedConfig, - models: &[crate::model::types::Model], - provider: &str, - model_id: &str, -) -> Option { - model_reasoning_capability(config, models, provider, model_id) - .and_then(|capability| capability.resolve(None)) -} - -pub(super) fn model_reasoning_capability( - config: &LoadedConfig, - models: &[crate::model::types::Model], - provider: &str, - model_id: &str, -) -> Option { - if let Some(capability) = models - .iter() - .find(|model| model.provider_id == provider && model.id == model_id) - .and_then(|model| { - crate::model::reasoning::capability_from_options(&model.reasoning_options) - }) - { - return Some(capability); - } - crate::model::discovery::Discovery::new_with_custom(Some( - config.merged_config.custom_providers.clone(), - )) - .ok() - .and_then(|discovery| discovery.get_model_reasoning_capability(provider, model_id)) - .filter(|capability| !capability.values().is_empty()) -} - -pub(super) fn model_context_window( - config: &LoadedConfig, - models: &[crate::model::types::Model], - provider: &str, - model: &str, -) -> Option { - if let Some(context_window) = models - .iter() - .find(|candidate| candidate.provider_id == provider && candidate.id == model) - .and_then(|model| model.context_window) - { - return Some(context_window); - } - - let discovery = crate::model::discovery::Discovery::new_with_custom(Some( - config.merged_config.custom_providers.clone(), - )) - .ok(); - discovery - .as_ref() - .and_then(|discovery| discovery.get_model_limit(provider, model)) - .or_else(|| { - config - .merged_config - .custom_providers - .get(provider) - .and_then(|provider| provider.models.get(model)) - .and_then(|model| model.context_window) - }) -} - impl AcpService { pub async fn set_mode( &self, @@ -390,211 +106,3 @@ impl AcpService { ))) } } - -#[cfg(test)] -mod tests { - use super::super::test_support::config_with_command; - use super::super::test_support::empty_config; - use super::super::test_support::model; - use super::super::test_support::test_session; - use super::*; - use agent_client_protocol::schema::v1::SessionConfigOptionCategory; - use std::path::PathBuf; - - #[test] - fn merges_client_mcp_servers_into_session_config() { - let mut config = config_with_command(crate::command::custom::CustomCommand { - name: "custom".to_string(), - description: None, - template: "$ARGUMENTS".to_string(), - agent: None, - model: None, - subtask: Some(false), - source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( - "/tmp/opencode.jsonc", - )), - workdir: PathBuf::from("/tmp"), - }); - merge_acp_mcp_servers( - &mut config, - vec![serde_json::from_value(serde_json::json!({ - "name": "zed-fs", - "command": "mcp-server", - "args": ["--stdio"], - "env": [{"name": "ROOT", "value": "/workspace"}] - })) - .expect("ACP stdio MCP server")], - ); - - let server = config - .merged_config - .mcp - .get("zed-fs") - .expect("merged ACP MCP server"); - assert_eq!(server.kind(), "local"); - assert!(server.enabled()); - } - - #[test] - fn builds_grouped_model_config_option() { - let option = model_config_option( - &[ - model("anthropic", "Anthropic", "claude", "Claude"), - model("openai", "OpenAI", "gpt-5", "GPT-5"), - model("openai", "OpenAI", "gpt-5-mini", "GPT-5 Mini"), - ], - "openai", - "gpt-5", - ); - - assert_eq!(option.id.to_string(), "model"); - assert_eq!(option.category, Some(SessionConfigOptionCategory::Model)); - let agent_client_protocol::schema::v1::SessionConfigKind::Select(select) = option.kind - else { - panic!("model option should be a select"); - }; - assert_eq!(select.current_value.to_string(), "openai/gpt-5"); - let agent_client_protocol::schema::v1::SessionConfigSelectOptions::Grouped(groups) = - select.options - else { - panic!("model options should be grouped"); - }; - assert_eq!(groups.len(), 2); - assert_eq!(groups[1].name, "OpenAI"); - assert_eq!(groups[1].options[1].value.to_string(), "openai/gpt-5-mini"); - } - - #[test] - fn validates_model_refs_against_selectable_catalog() { - let models = [model("openai", "OpenAI", "gpt-5", "GPT-5")]; - - assert_eq!( - find_selectable_model(&models, "openai/gpt-5") - .expect("known model") - .id, - "gpt-5" - ); - assert!(find_selectable_model(&models, "openai/missing").is_err()); - assert!(find_selectable_model(&models, "other/gpt-5").is_err()); - } - - #[test] - fn resolves_context_window_from_selectable_models() { - let mut model = model("example", "Example", "large-context", "Large Context"); - model.context_window = Some(1_090_000); - - assert_eq!( - model_context_window(&empty_config(), &[model], "example", "large-context"), - Some(1_090_000) - ); - } - - #[test] - fn preserves_selected_reasoning_effort_when_model_cannot_apply_it() { - let session = test_session(); - let option = reasoning_config_option(&session); - assert_eq!(option.id.to_string(), "effort"); - assert_eq!(option.name, "Effort"); - assert_eq!( - option.category, - Some(SessionConfigOptionCategory::ThoughtLevel) - ); - let agent_client_protocol::schema::v1::SessionConfigKind::Select(select) = option.kind - else { - panic!("reasoning option should be a select"); - }; - assert_eq!(select.current_value.to_string(), "high"); - let agent_client_protocol::schema::v1::SessionConfigSelectOptions::Ungrouped(options) = - select.options - else { - panic!("reasoning options should be ungrouped"); - }; - let values = options - .iter() - .map(|option| (option.value.to_string(), option.name.clone())) - .collect::>(); - assert_eq!( - values, - vec![ - ("none".to_string(), "None".to_string()), - ("low".to_string(), "Low".to_string()), - ("medium".to_string(), "Medium".to_string()), - ("high".to_string(), "High".to_string()), - ("xhigh".to_string(), "Xhigh".to_string()), - ("max".to_string(), "Max".to_string()), - ] - ); - } - - #[test] - fn compaction_mode_descriptions_preserve_configuration_and_visibility() { - let mut session = test_session(); - assert!(!session_modes(&session) - .available_modes - .iter() - .any(|mode| mode.id.to_string() == "compaction")); - - for description in ["", " \n ", "Use this agent to preserve project decisions."] { - let definitions = crate::agent::definition::parse_agent_definitions_from_config( - Some(&serde_json::json!({"compaction": { - "mode": "primary", "description": description - }})), - &mut Vec::new(), - ); - session.config.merged_config.agent_registry = - crate::agent::definition::AgentRegistry::with_definitions(None, definitions); - let expected = if description.trim().is_empty() { - "Compaction agent. Summarizes conversation context for continued work." - } else { - description - }; - let modes = session_modes(&session); - let compaction = modes - .available_modes - .iter() - .find(|mode| mode.id.to_string() == "compaction") - .expect("configured compaction mode"); - assert_eq!(compaction.description.as_deref(), Some(expected)); - let config = serde_json::to_value(session_config_options(&session)).unwrap(); - let compaction = config[0]["options"] - .as_array() - .unwrap() - .iter() - .find(|option| option["value"] == "compaction") - .expect("configured compaction option"); - assert_eq!(compaction["description"], expected); - assert_eq!( - modes - .available_modes - .iter() - .find(|mode| mode.id.to_string() == "build") - .unwrap() - .description - .as_deref(), - Some("The default agent. Executes tools based on configured permissions.") - ); - } - - for settings in [ - serde_json::json!({"mode": "subagent"}), - serde_json::json!({"mode": "primary", "hidden": true}), - ] { - let definitions = crate::agent::definition::parse_agent_definitions_from_config( - Some(&serde_json::json!({"compaction": settings})), - &mut Vec::new(), - ); - session.config.merged_config.agent_registry = - crate::agent::definition::AgentRegistry::with_definitions(None, definitions); - assert!(!session_modes(&session) - .available_modes - .iter() - .any(|mode| mode.id.to_string() == "compaction")); - let config = serde_json::to_value(session_config_options(&session)).unwrap(); - assert!(!config[0]["options"] - .as_array() - .unwrap() - .iter() - .any(|option| option["value"] == "compaction")); - } - } -} diff --git a/src/acp/service/context.rs b/src/acp/service/context.rs new file mode 100644 index 0000000..e17cbac --- /dev/null +++ b/src/acp/service/context.rs @@ -0,0 +1,197 @@ +pub(super) fn acp_context_tokens( + messages: &[crate::session::types::Message], + prompt_tokens: usize, +) -> usize { + let estimated = + crate::session::compaction::total_context_tokens(messages).saturating_add(prompt_tokens); + // Retained tail usage still describes the old context, as does the summary's billed usage. + let start = messages + .iter() + .rposition(crate::session::compaction::is_compaction_marker) + .map(|index| index + 1) + .unwrap_or(0); + for (index, message) in messages.iter().enumerate().skip(start).rev() { + if message.role != crate::session::types::MessageRole::Assistant { + continue; + } + // Each tool step bills its prompt again; cumulative usage is not a context size. + let provider_context = message + .parts + .iter() + .enumerate() + .filter_map(|(part_index, part)| { + let usage = part.recorded_usage()?; + let input = usage + .input + .saturating_add(usage.cache_read) + .saturating_add(usage.cache_write); + if input == 0 { + return None; + } + let mut tokens = + usize::try_from(input.saturating_add(usage.output)).unwrap_or(usize::MAX); + if let Some(previous_prompt) = part + .data + .get("acp_prompt_tokens") + .and_then(serde_json::Value::as_u64) + { + tokens = tokens + .saturating_sub(usize::try_from(previous_prompt).unwrap_or(usize::MAX)) + .saturating_add(prompt_tokens); + } + Some(tokens.saturating_add( + crate::session::compaction::assistant_parts_context_tokens( + &message.parts[part_index + 1..], + ), + )) + }) + .max() + .filter(|tokens| *tokens > 0); + if let Some(tokens) = provider_context { + let following = + crate::session::compaction::total_context_tokens(&messages[index + 1..]); + return estimated.max(tokens.saturating_add(following)); + } + } + estimated +} + +pub(super) fn estimate_acp_tokens(content: &str) -> usize { + content.chars().count().saturating_add(3) / 4 +} + +#[cfg(test)] +mod tests { + use super::super::transcript::record_tool_result; + use super::*; + + #[test] + fn acp_compaction_counts_provider_context_and_cached_input() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(60_000, 200, 20_000, 5_000, 0.0)); + let messages = vec![Message::user("question"), assistant]; + + assert_eq!(acp_context_tokens(&messages, 1_000), 85_200); + } + + #[test] + fn acp_compaction_does_not_sum_context_across_tool_requests() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(60_000, 200, 20_000, 0, 0.0)); + assistant + .parts + .push(MessagePart::usage(10_000, 100, 40_000, 0, 0.0)); + + assert_eq!( + acp_context_tokens(&[Message::user("question"), assistant], 0), + 80_200 + ); + } + + #[test] + fn acp_compaction_uses_latest_turn_usage_and_counts_pending_history() { + use crate::session::types::{Message, MessagePart}; + let mut old = Message::assistant("old answer"); + old.parts.push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); + let mut latest = Message::assistant("new answer"); + latest + .parts + .push(MessagePart::usage(50_000, 200, 0, 0, 0.0)); + let pending = Message::user("continue"); + let pending_tokens = crate::session::compaction::message_context_tokens(&pending); + let messages = vec![Message::user("question"), old, latest, pending]; + + assert_eq!(acp_context_tokens(&messages, 0), 50_200 + pending_tokens); + } + + #[test] + fn acp_compaction_counts_prompt_and_pending_tokens_without_provider_usage() { + use crate::config::configuration::CompactionConfig; + use crate::session::types::Message; + let messages = vec![Message::user("history ".repeat(2_000))]; + let history_tokens = crate::session::compaction::total_context_tokens(&messages); + let config = CompactionConfig::Settings { + auto: true, + prune: false, + reserved: Some(1_000), + }; + let context_window = Some(history_tokens as u32 + 3_000); + assert!(!crate::session::compaction::should_auto_compact( + &config, + acp_context_tokens(&messages, 1_500), + context_window, + None, + )); + assert!(crate::session::compaction::should_auto_compact( + &config, + acp_context_tokens(&messages, 1_500) + 500, + context_window, + None, + )); + } + + #[test] + fn acp_compaction_replaces_previous_prompt_estimate_in_provider_context() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + let mut usage = MessagePart::usage(80_000, 200, 0, 0, 0.0); + usage.data["acp_prompt_tokens"] = serde_json::json!(5_000); + assistant.parts.push(usage); + let messages = vec![Message::user("question"), assistant]; + assert_eq!(acp_context_tokens(&messages, 25_000), 100_200); + assert_eq!(acp_context_tokens(&messages, 1_000), 76_200); + } + + #[test] + fn acp_compaction_keeps_provider_context_when_new_input_usage_is_missing() { + use crate::session::types::{Message, MessagePart}; + for output in [0, 100] { + let mut previous = Message::assistant("previous answer"); + previous + .parts + .push(MessagePart::usage(85_000, 200, 0, 0, 0.0)); + let mut latest = Message::assistant("new answer"); + let mut usage = MessagePart::usage(0, output, 0, 0, 0.0); + usage.data["acp_prompt_tokens"] = serde_json::json!(1_000); + latest.parts.push(usage); + let following_tokens = crate::session::compaction::message_context_tokens(&latest); + assert_eq!( + acp_context_tokens(&[Message::user("question"), previous, latest], 2_000), + 85_200 + following_tokens + ); + } + } + + #[test] + fn acp_compaction_counts_tool_output_after_provider_usage() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(80_000, 200, 0, 0, 0.0)); + record_tool_result( + &mut assistant, + &crate::llm::ToolCallResult { + tool_call_id: "call_1".into(), + role: "tool".into(), + name: "read".into(), + content: serde_json::json!({"output_preview": "late tool output ".repeat(1_000)}) + .to_string(), + }, + ); + let mut following = Message::assistant(""); + following.parts = assistant.parts[2..].to_vec(); + let extra_tokens = crate::session::compaction::message_context_tokens(&following); + assert!(extra_tokens > 1_000); + assert_eq!( + acp_context_tokens(&[Message::user("question"), assistant], 0), + 80_200 + extra_tokens + ); + } +} diff --git a/src/acp/service/mcp_config.rs b/src/acp/service/mcp_config.rs new file mode 100644 index 0000000..2aafc7d --- /dev/null +++ b/src/acp/service/mcp_config.rs @@ -0,0 +1,112 @@ +use crate::config::configuration::LoadedConfig; +use agent_client_protocol::schema::v1::McpServer; + +pub(super) fn merge_acp_mcp_servers(config: &mut LoadedConfig, servers: Vec) { + for server in servers { + let (name, server) = match server { + McpServer::Stdio(server) => { + let mut command = vec![server.command.to_string_lossy().into_owned()]; + command.extend(server.args); + ( + server.name, + crate::config::configuration::McpServerConfig::Local( + crate::config::configuration::McpLocalConfig { + command, + cwd: None, + environment: server + .env + .into_iter() + .map(|variable| (variable.name, variable.value)) + .collect(), + enabled: true, + timeout_ms: None, + }, + ), + ) + } + McpServer::Http(server) => ( + server.name, + crate::config::configuration::McpServerConfig::Remote( + crate::config::configuration::McpRemoteConfig { + url: server.url, + headers: server + .headers + .into_iter() + .map(|header| (header.name, header.value)) + .collect(), + enabled: true, + timeout_ms: None, + oauth_enabled: false, + oauth_client_id: None, + oauth_client_secret: None, + oauth_scope: None, + }, + ), + ), + McpServer::Sse(server) => ( + server.name, + crate::config::configuration::McpServerConfig::Remote( + crate::config::configuration::McpRemoteConfig { + url: server.url, + headers: server + .headers + .into_iter() + .map(|header| (header.name, header.value)) + .collect(), + enabled: true, + timeout_ms: None, + oauth_enabled: false, + oauth_client_id: None, + oauth_client_secret: None, + oauth_scope: None, + }, + ), + ), + #[allow(unreachable_patterns)] + _ => continue, + }; + config.merged_config.mcp.insert(name, server); + } +} + +#[cfg(test)] +mod tests { + use super::super::test_support::config_with_command; + + use super::*; + use std::path::PathBuf; + + #[test] + fn merges_client_mcp_servers_into_session_config() { + let mut config = config_with_command(crate::command::custom::CustomCommand { + name: "custom".to_string(), + description: None, + template: "$ARGUMENTS".to_string(), + agent: None, + model: None, + subtask: Some(false), + source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( + "/tmp/opencode.jsonc", + )), + workdir: PathBuf::from("/tmp"), + }); + merge_acp_mcp_servers( + &mut config, + vec![serde_json::from_value(serde_json::json!({ + "name": "zed-fs", + "command": "mcp-server", + "args": ["--stdio"], + "env": [{"name": "ROOT", "value": "/workspace"}] + })) + .expect("ACP stdio MCP server")], + ); + + let server = config + .merged_config + .mcp + .get("zed-fs") + .expect("merged ACP MCP server"); + assert_eq!(server.kind(), "local"); + assert!(server.enabled()); + } +} diff --git a/src/acp/service/mod.rs b/src/acp/service/mod.rs index f121f7a..aa38437 100644 --- a/src/acp/service/mod.rs +++ b/src/acp/service/mod.rs @@ -1,14 +1,26 @@ +mod attachment_storage; mod attachments; +mod command_expansion; mod commands; mod compaction; +mod compaction_summary; mod config; +mod context; +mod mcp_config; +mod models; +mod modes; mod notifications; +mod permission_preview; mod permissions; mod prompt; mod questions; +mod replay; mod sessions; mod stream; mod terminals; +mod tool_content; +mod tool_locations; +mod transcript; mod usage; #[cfg(test)] diff --git a/src/acp/service/models.rs b/src/acp/service/models.rs new file mode 100644 index 0000000..55464b9 --- /dev/null +++ b/src/acp/service/models.rs @@ -0,0 +1,300 @@ +use super::AcpSession; +use crate::config::configuration::LoadedConfig; +use agent_client_protocol::schema::v1::SessionConfigOption; +use agent_client_protocol::schema::v1::SessionConfigOptionCategory; +use agent_client_protocol::schema::v1::SessionConfigSelectGroup; +use agent_client_protocol::schema::v1::SessionConfigSelectOption; +use agent_client_protocol::Error; + +pub(super) fn model_output_limit( + config: &LoadedConfig, + provider_id: &str, + model_id: &str, +) -> Option { + if let Some(limit) = config + .merged_config + .custom_providers + .get(&provider_id.trim().to_ascii_lowercase()) + .and_then(|provider| provider.models.get(model_id)) + .and_then(|model| model.max_tokens) + { + return Some(limit); + } + crate::model::discovery::Discovery::new_with_config( + Some(config.merged_config.custom_providers.clone()), + config.merged_config.disabled_providers.clone(), + config.merged_config.enabled_providers.clone(), + ) + .ok() + .and_then(|discovery| discovery.get_model_output_limit(provider_id, model_id)) +} + +pub(super) fn resolved_reasoning( + session: &AcpSession, + requested: crate::model::reasoning::ReasoningEffort, +) -> Option { + model_reasoning_capability( + &session.config, + &session.models, + &session.provider, + &session.model, + ) + .and_then(|capability| capability.resolve(Some(requested))) +} + +fn reasoning_effort_label(effort: crate::model::reasoning::ReasoningEffort) -> &'static str { + match effort { + crate::model::reasoning::ReasoningEffort::None => "None", + crate::model::reasoning::ReasoningEffort::Minimal => "Minimal", + crate::model::reasoning::ReasoningEffort::Low => "Low", + crate::model::reasoning::ReasoningEffort::Medium => "Medium", + crate::model::reasoning::ReasoningEffort::High => "High", + crate::model::reasoning::ReasoningEffort::XHigh => "Xhigh", + crate::model::reasoning::ReasoningEffort::Max => "Max", + } +} + +pub(super) fn resolve_model(config: &LoadedConfig) -> (String, String) { + if let Ok(Some((provider, model))) = + crate::persistence::PrefsDAO::new().and_then(|prefs| prefs.get_active_model()) + { + return (provider, model); + } + config + .merged_config + .model + .as_deref() + .map(crate::app::parse_model_ref) + .unwrap_or_else(|| ("opencode".to_string(), "big-pickle".to_string())) +} + +pub(super) fn reasoning_config_option(session: &AcpSession) -> SessionConfigOption { + let options = [ + crate::model::reasoning::ReasoningEffort::None, + crate::model::reasoning::ReasoningEffort::Low, + crate::model::reasoning::ReasoningEffort::Medium, + crate::model::reasoning::ReasoningEffort::High, + crate::model::reasoning::ReasoningEffort::XHigh, + crate::model::reasoning::ReasoningEffort::Max, + ] + .iter() + .map(|effort| SessionConfigSelectOption::new(effort.as_str(), reasoning_effort_label(*effort))) + .collect::>(); + let current = session.reasoning_selection.as_str(); + SessionConfigOption::select("effort", "Effort", current, options) + .category(SessionConfigOptionCategory::ThoughtLevel) +} + +pub(super) fn model_config_option( + models: &[crate::model::types::Model], + provider: &str, + model_id: &str, +) -> SessionConfigOption { + let mut model_groups: Vec = Vec::new(); + for model in models { + let option = SessionConfigSelectOption::new(model_value(model), model.name.clone()) + .description(model.id.clone()); + if model_groups + .last() + .is_some_and(|group| group.group.to_string() == model.provider_id) + { + model_groups + .last_mut() + .expect("model group exists") + .options + .push(option); + } else { + model_groups.push(SessionConfigSelectGroup::new( + model.provider_id.clone(), + model.provider_name.clone(), + vec![option], + )); + } + } + + SessionConfigOption::select( + "model", + "Model", + format!("{provider}/{model_id}"), + model_groups, + ) + .category(SessionConfigOptionCategory::Model) +} + +fn model_value(model: &crate::model::types::Model) -> String { + crate::model::catalog::model_ref(model) +} + +pub(super) fn find_selectable_model<'a>( + models: &'a [crate::model::types::Model], + model_ref: &str, +) -> Result<&'a crate::model::types::Model, Error> { + models + .iter() + .find(|model| model_value(model) == model_ref) + .ok_or_else(|| Error::invalid_params().data("unknown or unavailable ACP model")) +} + +pub(super) fn model_reasoning( + config: &LoadedConfig, + models: &[crate::model::types::Model], + provider: &str, + model_id: &str, +) -> Option { + model_reasoning_capability(config, models, provider, model_id) + .and_then(|capability| capability.resolve(None)) +} + +pub(super) fn model_reasoning_capability( + config: &LoadedConfig, + models: &[crate::model::types::Model], + provider: &str, + model_id: &str, +) -> Option { + if let Some(capability) = models + .iter() + .find(|model| model.provider_id == provider && model.id == model_id) + .and_then(|model| { + crate::model::reasoning::capability_from_options(&model.reasoning_options) + }) + { + return Some(capability); + } + crate::model::discovery::Discovery::new_with_custom(Some( + config.merged_config.custom_providers.clone(), + )) + .ok() + .and_then(|discovery| discovery.get_model_reasoning_capability(provider, model_id)) + .filter(|capability| !capability.values().is_empty()) +} + +pub(super) fn model_context_window( + config: &LoadedConfig, + models: &[crate::model::types::Model], + provider: &str, + model: &str, +) -> Option { + if let Some(context_window) = models + .iter() + .find(|candidate| candidate.provider_id == provider && candidate.id == model) + .and_then(|model| model.context_window) + { + return Some(context_window); + } + + let discovery = crate::model::discovery::Discovery::new_with_custom(Some( + config.merged_config.custom_providers.clone(), + )) + .ok(); + discovery + .as_ref() + .and_then(|discovery| discovery.get_model_limit(provider, model)) + .or_else(|| { + config + .merged_config + .custom_providers + .get(provider) + .and_then(|provider| provider.models.get(model)) + .and_then(|model| model.context_window) + }) +} + +#[cfg(test)] +mod tests { + use super::super::test_support::empty_config; + use super::super::test_support::model; + use super::super::test_support::test_session; + use super::*; + use agent_client_protocol::schema::v1::SessionConfigOptionCategory; + + #[test] + fn builds_grouped_model_config_option() { + let option = model_config_option( + &[ + model("anthropic", "Anthropic", "claude", "Claude"), + model("openai", "OpenAI", "gpt-5", "GPT-5"), + model("openai", "OpenAI", "gpt-5-mini", "GPT-5 Mini"), + ], + "openai", + "gpt-5", + ); + + assert_eq!(option.id.to_string(), "model"); + assert_eq!(option.category, Some(SessionConfigOptionCategory::Model)); + let agent_client_protocol::schema::v1::SessionConfigKind::Select(select) = option.kind + else { + panic!("model option should be a select"); + }; + assert_eq!(select.current_value.to_string(), "openai/gpt-5"); + let agent_client_protocol::schema::v1::SessionConfigSelectOptions::Grouped(groups) = + select.options + else { + panic!("model options should be grouped"); + }; + assert_eq!(groups.len(), 2); + assert_eq!(groups[1].name, "OpenAI"); + assert_eq!(groups[1].options[1].value.to_string(), "openai/gpt-5-mini"); + } + + #[test] + fn validates_model_refs_against_selectable_catalog() { + let models = [model("openai", "OpenAI", "gpt-5", "GPT-5")]; + + assert_eq!( + find_selectable_model(&models, "openai/gpt-5") + .expect("known model") + .id, + "gpt-5" + ); + assert!(find_selectable_model(&models, "openai/missing").is_err()); + assert!(find_selectable_model(&models, "other/gpt-5").is_err()); + } + + #[test] + fn resolves_context_window_from_selectable_models() { + let mut model = model("example", "Example", "large-context", "Large Context"); + model.context_window = Some(1_090_000); + + assert_eq!( + model_context_window(&empty_config(), &[model], "example", "large-context"), + Some(1_090_000) + ); + } + + #[test] + fn preserves_selected_reasoning_effort_when_model_cannot_apply_it() { + let session = test_session(); + let option = reasoning_config_option(&session); + assert_eq!(option.id.to_string(), "effort"); + assert_eq!(option.name, "Effort"); + assert_eq!( + option.category, + Some(SessionConfigOptionCategory::ThoughtLevel) + ); + let agent_client_protocol::schema::v1::SessionConfigKind::Select(select) = option.kind + else { + panic!("reasoning option should be a select"); + }; + assert_eq!(select.current_value.to_string(), "high"); + let agent_client_protocol::schema::v1::SessionConfigSelectOptions::Ungrouped(options) = + select.options + else { + panic!("reasoning options should be ungrouped"); + }; + let values = options + .iter() + .map(|option| (option.value.to_string(), option.name.clone())) + .collect::>(); + assert_eq!( + values, + vec![ + ("none".to_string(), "None".to_string()), + ("low".to_string(), "Low".to_string()), + ("medium".to_string(), "Medium".to_string()), + ("high".to_string(), "High".to_string()), + ("xhigh".to_string(), "Xhigh".to_string()), + ("max".to_string(), "Max".to_string()), + ] + ); + } +} diff --git a/src/acp/service/modes.rs b/src/acp/service/modes.rs new file mode 100644 index 0000000..82634e6 --- /dev/null +++ b/src/acp/service/modes.rs @@ -0,0 +1,106 @@ +use super::AcpSession; +use agent_client_protocol::schema::v1::SessionMode; +use agent_client_protocol::schema::v1::SessionModeState; + +pub(super) fn mode_description(agent: &crate::agent::definition::AgentDefinition) -> &str { + if agent.name == "compaction" && agent.description.trim().is_empty() { + "Compaction agent. Summarizes conversation context for continued work." + } else { + &agent.description + } +} + +pub(super) fn session_modes(session: &AcpSession) -> SessionModeState { + let mut modes = session + .config + .merged_config + .agent_registry + .visible_primary_agents() + .into_iter() + .map(|agent| { + SessionMode::new(agent.name.clone(), agent.name.clone()) + .description(mode_description(agent)) + }) + .collect::>(); + modes.sort_by(|left, right| left.name.cmp(&right.name)); + SessionModeState::new(session.agent.clone(), modes) +} + +#[cfg(test)] +mod tests { + use super::super::config::session_config_options; + use super::super::test_support::test_session; + use super::*; + + #[test] + fn compaction_mode_descriptions_preserve_configuration_and_visibility() { + let mut session = test_session(); + assert!(!session_modes(&session) + .available_modes + .iter() + .any(|mode| mode.id.to_string() == "compaction")); + + for description in ["", " \n ", "Use this agent to preserve project decisions."] { + let definitions = crate::agent::definition::parse_agent_definitions_from_config( + Some(&serde_json::json!({"compaction": { + "mode": "primary", "description": description + }})), + &mut Vec::new(), + ); + session.config.merged_config.agent_registry = + crate::agent::definition::AgentRegistry::with_definitions(None, definitions); + let expected = if description.trim().is_empty() { + "Compaction agent. Summarizes conversation context for continued work." + } else { + description + }; + let modes = session_modes(&session); + let compaction = modes + .available_modes + .iter() + .find(|mode| mode.id.to_string() == "compaction") + .expect("configured compaction mode"); + assert_eq!(compaction.description.as_deref(), Some(expected)); + let config = serde_json::to_value(session_config_options(&session)).unwrap(); + let compaction = config[0]["options"] + .as_array() + .unwrap() + .iter() + .find(|option| option["value"] == "compaction") + .expect("configured compaction option"); + assert_eq!(compaction["description"], expected); + assert_eq!( + modes + .available_modes + .iter() + .find(|mode| mode.id.to_string() == "build") + .unwrap() + .description + .as_deref(), + Some("The default agent. Executes tools based on configured permissions.") + ); + } + + for settings in [ + serde_json::json!({"mode": "subagent"}), + serde_json::json!({"mode": "primary", "hidden": true}), + ] { + let definitions = crate::agent::definition::parse_agent_definitions_from_config( + Some(&serde_json::json!({"compaction": settings})), + &mut Vec::new(), + ); + session.config.merged_config.agent_registry = + crate::agent::definition::AgentRegistry::with_definitions(None, definitions); + assert!(!session_modes(&session) + .available_modes + .iter() + .any(|mode| mode.id.to_string() == "compaction")); + let config = serde_json::to_value(session_config_options(&session)).unwrap(); + assert!(!config[0]["options"] + .as_array() + .unwrap() + .iter() + .any(|option| option["value"] == "compaction")); + } + } +} diff --git a/src/acp/service/notifications.rs b/src/acp/service/notifications.rs index ef4ec5d..ffc2c64 100644 --- a/src/acp/service/notifications.rs +++ b/src/acp/service/notifications.rs @@ -1,11 +1,11 @@ use super::internal_error; -use agent_client_protocol::schema::v1::ContentBlock; +use super::tool_content::tool_result_content; +use super::tool_locations::tool_locations; +use super::tool_locations::tool_result_locations; use agent_client_protocol::schema::v1::ContentChunk; use agent_client_protocol::schema::v1::SessionNotification; use agent_client_protocol::schema::v1::SessionUpdate; use agent_client_protocol::schema::v1::ToolCall; -use agent_client_protocol::schema::v1::ToolCallContent; -use agent_client_protocol::schema::v1::ToolCallLocation; use agent_client_protocol::schema::v1::ToolCallStatus; use agent_client_protocol::schema::v1::ToolCallUpdate; use agent_client_protocol::schema::v1::ToolCallUpdateFields; @@ -14,49 +14,6 @@ use agent_client_protocol::Client; use agent_client_protocol::ConnectionTo; use agent_client_protocol::Error; use std::path::Path; -use std::path::PathBuf; - -fn mcp_native_tool_content(metadata: &serde_json::Value) -> Vec { - metadata - .get("mcp_result") - .and_then(|result| result.get("content")) - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter(|block| { - !matches!( - block.get("type").and_then(serde_json::Value::as_str), - Some("text") - ) - }) - .map(|block| { - serde_json::from_value::(block.clone()) - .map(ToolCallContent::from) - .unwrap_or_else(|_| { - let kind = block - .get("type") - .and_then(serde_json::Value::as_str) - .unwrap_or("unknown"); - ToolCallContent::from(format!( - "[MCP {kind} content]\n{}", - serde_json::to_string_pretty(block).unwrap_or_else(|_| block.to_string()) - )) - }) - }) - .collect() -} - -fn metadata_has_mcp_images(metadata: Option<&serde_json::Value>) -> bool { - metadata - .and_then(|metadata| metadata.get("mcp_result")) - .and_then(|result| result.get("content")) - .and_then(serde_json::Value::as_array) - .is_some_and(|content| { - content - .iter() - .any(|block| block.get("type").and_then(serde_json::Value::as_str) == Some("image")) - }) -} pub(super) fn send_text( connection: &ConnectionTo, @@ -75,141 +32,6 @@ pub(super) fn send_text( .map_err(|_| internal_error()) } -fn tool_result_text(payload: &serde_json::Value) -> String { - payload - .get("output") - .or_else(|| payload.get("output_preview")) - .or_else(|| payload.get("error")) - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string() -} - -pub(super) fn absolute_tool_path(path: &str, cwd: &Path) -> PathBuf { - let path = PathBuf::from(path); - if path.is_absolute() { - path - } else { - cwd.join(path) - } -} - -pub(super) fn tool_locations( - tool_name: &str, - input: &serde_json::Value, - cwd: &Path, -) -> Vec { - let paths = if tool_name == "write_files" { - input - .get("files") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(|file| file.get("file_path").and_then(serde_json::Value::as_str)) - .map(str::to_string) - .collect() - } else if tool_name == "apply_patch" { - crate::tools::patch::patch_paths_from_params(input) - } else { - input - .get("file_path") - .or_else(|| input.get("filePath")) - .or_else(|| input.get("filepath")) - .or_else(|| input.get("path")) - .and_then(serde_json::Value::as_str) - .map(|path| vec![path.to_string()]) - .unwrap_or_default() - }; - - paths - .into_iter() - .map(|path| ToolCallLocation::new(absolute_tool_path(&path, cwd))) - .collect() -} - -fn tool_result_locations(payload: &serde_json::Value, cwd: &Path) -> Vec { - let Some(metadata) = payload.get("metadata") else { - return Vec::new(); - }; - let line = metadata - .get("line_number") - .and_then(serde_json::Value::as_u64) - .and_then(|line| u32::try_from(line).ok()); - - if let Some(changes) = metadata - .get("changes") - .and_then(serde_json::Value::as_array) - { - return changes - .iter() - .filter_map(|change| change.get("path").and_then(serde_json::Value::as_str)) - .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd))) - .collect(); - } - - metadata - .get("path") - .and_then(serde_json::Value::as_str) - .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd)).line(line)) - .into_iter() - .collect() -} - -fn tool_result_content(payload: &serde_json::Value, cwd: &Path) -> Vec { - let mut content = Vec::new(); - let text = tool_result_text(payload); - if !text.is_empty() { - content.push(ToolCallContent::from(text)); - } - - if let Some(metadata) = payload.get("metadata") { - if let Some(changes) = metadata - .get("changes") - .and_then(serde_json::Value::as_array) - { - content.extend(changes.iter().filter_map(|change| tool_diff(change, cwd))); - } else if let Some(diff) = tool_diff(metadata, cwd) { - content.push(diff); - } - content.extend(mcp_native_tool_content(metadata)); - } - - if !metadata_has_mcp_images(payload.get("metadata")) { - if let Some(images) = payload.get("images").and_then(serde_json::Value::as_array) { - content.extend(images.iter().filter_map(tool_result_image)); - } - } - - content -} - -fn tool_diff(change: &serde_json::Value, cwd: &Path) -> Option { - let path = change.get("path")?.as_str()?; - let new_text = change.get("new_text")?.as_str()?; - let old_text = change.get("old_text").and_then(serde_json::Value::as_str); - Some( - agent_client_protocol::schema::v1::Diff::new( - absolute_tool_path(path, cwd), - new_text.to_string(), - ) - .old_text(old_text.map(str::to_string)) - .into(), - ) -} - -fn tool_result_image(image: &serde_json::Value) -> Option { - let data = image.get("data_url")?.as_str()?; - let media_type = image.get("media_type")?.as_str()?; - let encoded = data - .strip_prefix("data:") - .and_then(|value| value.split_once(',')) - .map(|(_, encoded)| encoded) - .unwrap_or(data); - Some(ToolCallContent::from(ContentBlock::Image( - agent_client_protocol::schema::v1::ImageContent::new(encoded, media_type), - ))) -} - pub(super) fn send_tool_call( connection: &ConnectionTo, session_id: &str, @@ -259,152 +81,6 @@ pub(super) fn send_tool_result( .map_err(|_| internal_error()) } -pub(super) fn replay_messages( - connection: &ConnectionTo, - session_id: &str, - messages: &[crate::session::types::Message], - cwd: &Path, -) -> Result<(), Error> { - for message in messages { - if crate::session::compaction::is_compaction_display_item(message) { - continue; - } - let message_id = message.id.clone(); - match message.role { - crate::session::types::MessageRole::User => { - if !message.content.is_empty() { - send_replay_text( - connection, - session_id, - &message_id, - &message.content, - true, - false, - )?; - } - } - crate::session::types::MessageRole::Assistant => { - for part in &message.parts { - match part.part_type.as_str() { - "text" => { - if let Some(text) = part.text_value() { - send_replay_text( - connection, - session_id, - &message_id, - text, - false, - false, - )?; - } - } - "reasoning" => { - if let Some(text) = part.text_value() { - send_replay_text( - connection, - session_id, - &message_id, - text, - false, - true, - )?; - } - } - "tool_call" => replay_tool_call(connection, session_id, part, cwd)?, - "tool_result" => replay_tool_result(connection, session_id, part, cwd)?, - _ => {} - } - } - } - crate::session::types::MessageRole::System - | crate::session::types::MessageRole::Tool => {} - } - } - Ok(()) -} - -pub(super) fn send_replay_text( - connection: &ConnectionTo, - session_id: &str, - message_id: &str, - text: &str, - user: bool, - thought: bool, -) -> Result<(), Error> { - let chunk = ContentChunk::new(text.to_string().into()).message_id(message_id); - let update = if user { - SessionUpdate::UserMessageChunk(chunk) - } else if thought { - SessionUpdate::AgentThoughtChunk(chunk) - } else { - SessionUpdate::AgentMessageChunk(chunk) - }; - connection - .send_notification(SessionNotification::new(session_id.to_string(), update)) - .map_err(|_| internal_error()) -} - -fn replay_tool_call( - connection: &ConnectionTo, - session_id: &str, - part: &crate::session::types::MessagePart, - cwd: &Path, -) -> Result<(), Error> { - let Some(tool_call_id) = part.tool_id() else { - return Ok(()); - }; - let name = part.tool_name().unwrap_or("tool"); - let input = part - .data - .get("args") - .cloned() - .unwrap_or(serde_json::Value::Null); - let status = match part.tool_status() { - Some("completed") | Some("ok") => ToolCallStatus::Completed, - Some("error") | Some("failed") => ToolCallStatus::Failed, - Some("running") => ToolCallStatus::InProgress, - _ => ToolCallStatus::Pending, - }; - let update = SessionUpdate::ToolCall( - ToolCall::new(tool_call_id.to_string(), tool_title(name, &input)) - .kind(tool_kind(name)) - .status(status) - .locations(tool_locations(name, &input, cwd)) - .raw_input(input), - ); - connection - .send_notification(SessionNotification::new(session_id.to_string(), update)) - .map_err(|_| internal_error()) -} - -fn replay_tool_result( - connection: &ConnectionTo, - session_id: &str, - part: &crate::session::types::MessagePart, - cwd: &Path, -) -> Result<(), Error> { - let Some(tool_call_id) = part.tool_id() else { - return Ok(()); - }; - let content = part - .data - .get("content") - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string(); - send_tool_result( - connection, - session_id, - crate::llm::ToolCallResult { - tool_call_id: tool_call_id.to_string(), - role: "tool".to_string(), - name: part.tool_name().unwrap_or("tool").to_string(), - content, - }, - cwd, - ) -} - pub(super) fn tool_kind(tool_name: &str) -> ToolKind { match tool_name { "bash" | "bash_output" | "bash_kill" | "bash_restart" | "terminal_session" => { @@ -419,7 +95,7 @@ pub(super) fn tool_kind(tool_name: &str) -> ToolKind { } } -fn tool_title(tool_name: &str, input: &serde_json::Value) -> String { +pub(super) fn tool_title(tool_name: &str, input: &serde_json::Value) -> String { match tool_name { "bash" => input .get("command") @@ -441,11 +117,7 @@ fn tool_title(tool_name: &str, input: &serde_json::Value) -> String { #[cfg(test)] mod tests { use super::*; - use agent_client_protocol::schema::v1::ContentBlock; - use agent_client_protocol::schema::v1::ToolCallContent; use agent_client_protocol::schema::v1::ToolKind; - use std::path::Path; - use std::path::PathBuf; #[test] fn maps_crabcode_tools_to_acp_kinds() { @@ -466,157 +138,4 @@ mod tests { "src/main.rs" ); } - - #[test] - fn acp_tool_result_prefers_full_output() { - let payload = serde_json::json!({ - "output": "complete tool output", - "output_preview": "short preview", - }); - - assert_eq!(tool_result_text(&payload), "complete tool output"); - } - - #[test] - fn acp_tool_result_supports_legacy_preview_payloads() { - let payload = serde_json::json!({"output_preview": "legacy output"}); - - assert_eq!(tool_result_text(&payload), "legacy output"); - } - - #[test] - fn acp_tool_locations_normalize_multi_file_and_patch_paths() { - let cwd = Path::new("/tmp/workspace"); - let write_locations = tool_locations( - "write_files", - &serde_json::json!({ - "files": [ - {"file_path": "src/a.rs", "content": "a"}, - {"file_path": "/tmp/b.rs", "content": "b"} - ] - }), - cwd, - ); - assert_eq!( - write_locations[0].path, - PathBuf::from("/tmp/workspace/src/a.rs") - ); - assert_eq!(write_locations[1].path, PathBuf::from("/tmp/b.rs")); - - let patch_locations = tool_locations( - "apply_patch", - &serde_json::json!({ - "patch": "*** Begin Patch\n*** Update File: src/a.rs\n*** Add File: src/b.rs\n*** End Patch" - }), - cwd, - ); - assert_eq!(patch_locations.len(), 2); - assert_eq!( - patch_locations[1].path, - PathBuf::from("/tmp/workspace/src/b.rs") - ); - } - - #[test] - fn acp_tool_result_emits_diff_location_and_image_content() { - let payload = serde_json::json!({ - "output": "updated", - "metadata": { - "path": "src/main.rs", - "line_number": 4, - "old_text": "fn old() {}", - "new_text": "fn new() {}" - }, - "images": [{ - "data_url": "data:image/png;base64,aGk=", - "media_type": "image/png" - }] - }); - let cwd = Path::new("/tmp/workspace"); - - let locations = tool_result_locations(&payload, cwd); - assert_eq!( - locations[0].path, - PathBuf::from("/tmp/workspace/src/main.rs") - ); - assert_eq!(locations[0].line, Some(4)); - - let content = tool_result_content(&payload, cwd); - let diff = content.iter().find_map(|item| match item { - ToolCallContent::Diff(diff) => Some(diff), - _ => None, - }); - let diff = diff.expect("diff content"); - assert_eq!(diff.path, PathBuf::from("/tmp/workspace/src/main.rs")); - assert_eq!(diff.old_text.as_deref(), Some("fn old() {}")); - assert_eq!(diff.new_text, "fn new() {}"); - assert!(content.iter().any(|item| matches!( - item, - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::Image(image) if image.data == "aGk=" && image.mime_type == "image/png") - ))); - } - - #[test] - fn acp_tool_result_restores_native_mcp_resources() { - let metadata = serde_json::json!({ - "mcp_result": { - "content": [ - { - "type": "image", - "data": "aGk=", - "mimeType": "image/png", - "annotations": { "priority": 0.5 } - }, - { - "type": "resource_link", - "uri": "file:///tmp/readme.md", - "name": "readme", - "mimeType": "text/markdown", - "annotations": { "priority": 0.8 } - } - ] - } - }); - - let content = mcp_native_tool_content(&metadata); - assert_eq!(content.len(), 2); - assert!(matches!( - &content[0], - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::Image(image) - if image.data == "aGk=" && image.annotations.is_some()) - )); - assert!(matches!( - &content[1], - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::ResourceLink(link) - if link.uri == "file:///tmp/readme.md") - )); - } - - #[test] - fn acp_tool_result_restores_mcp_audio_and_preserves_unknown_blocks() { - let metadata = serde_json::json!({ - "mcp_result": { - "content": [ - { "type": "audio", "data": "YXVkaW8=", "mimeType": "audio/wav" }, - { "type": "future_media", "payload": { "value": 1 } } - ] - } - }); - let content = mcp_native_tool_content(&metadata); - assert!(matches!( - &content[0], - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::Audio(audio) - if audio.data == "YXVkaW8=" && audio.mime_type == "audio/wav") - )); - assert!(matches!( - &content[1], - ToolCallContent::Content(content) - if matches!(&content.content, ContentBlock::Text(text) - if text.text.contains("future_media") && text.text.contains("payload")) - )); - } } diff --git a/src/acp/service/permission_preview.rs b/src/acp/service/permission_preview.rs new file mode 100644 index 0000000..361d8bb --- /dev/null +++ b/src/acp/service/permission_preview.rs @@ -0,0 +1,226 @@ +use super::tool_locations::absolute_tool_path; +use agent_client_protocol::schema::v1::ToolCallContent; +use std::path::Path; + +pub(super) fn permission_tool_content( + prompt: &crate::tools::PermissionPrompt, + cwd: &Path, +) -> Vec { + match prompt.tool_id.as_str() { + "write" => prompt + .raw_input + .get("file_path") + .or_else(|| prompt.raw_input.get("filePath")) + .and_then(serde_json::Value::as_str) + .zip( + prompt + .raw_input + .get("content") + .and_then(serde_json::Value::as_str), + ) + .map(|(path, new_text)| vec![preflight_diff(path, new_text, cwd)]) + .unwrap_or_default(), + "write_files" => prompt + .raw_input + .get("files") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|file| { + let path = file.get("file_path")?.as_str()?; + let content = file.get("content")?.as_str()?; + Some(preflight_diff(path, content, cwd)) + }) + .collect(), + "edit" => permission_edit_diff(&prompt.raw_input, cwd) + .into_iter() + .collect(), + "apply_patch" + if crate::tools::patch::patch_paths_from_params(&prompt.raw_input) + .iter() + .all(|path| preview_read_allowed(&absolute_tool_path(path, cwd), cwd)) => + { + crate::tools::patch::preview_patch(&prompt.raw_input, cwd) + .unwrap_or_default() + .into_iter() + .map(|change| { + agent_client_protocol::schema::v1::Diff::new(change.path, change.new_text) + .old_text(change.old_text) + .into() + }) + .collect() + } + _ => Vec::new(), + } +} + +// A write/edit approval is not a read approval. Keep ordinary workspace previews, +// but never acquire protected file contents just to display a permission request. +fn preview_read_allowed(path: &Path, cwd: &Path) -> bool { + !crate::tools::permission::is_outside_workdir(path, cwd) + && !crate::tools::permission::is_sensitive_path(path) + && !std::fs::canonicalize(path) + .ok() + .is_some_and(|resolved| crate::tools::permission::is_sensitive_path(&resolved)) +} + +fn preflight_diff(path: &str, new_text: &str, cwd: &Path) -> ToolCallContent { + let path = absolute_tool_path(path, cwd); + let old_text = preview_read_allowed(&path, cwd) + .then(|| std::fs::read_to_string(&path).ok()) + .flatten(); + agent_client_protocol::schema::v1::Diff::new(path, new_text.to_string()) + .old_text(old_text) + .into() +} + +fn permission_edit_diff(input: &serde_json::Value, cwd: &Path) -> Option { + let path = input + .get("file_path") + .or_else(|| input.get("filePath"))? + .as_str()?; + let old_string = input.get("old_string")?.as_str()?; + let new_string = input.get("new_string")?.as_str()?; + let absolute = absolute_tool_path(path, cwd); + if !preview_read_allowed(&absolute, cwd) { + return None; + } + let old_text = std::fs::read_to_string(&absolute).ok()?; + let new_text = if input + .get("replace_all") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) + { + old_text.replace(old_string, new_string) + } else { + old_text.replacen(old_string, new_string, 1) + }; + Some( + agent_client_protocol::schema::v1::Diff::new(absolute, new_text) + .old_text(old_text) + .into(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use agent_client_protocol::schema::v1::ToolCallContent; + + #[test] + fn acp_permission_edit_includes_preflight_diff() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("main.rs"); + std::fs::write(&path, "fn old() {}\n").unwrap(); + let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); + let prompt = crate::tools::PermissionPrompt { + tool_call_id: Some("call_edit".to_string()), + tool_id: "edit".to_string(), + action: crate::tools::PermissionAction::Write, + permission: "edit".to_string(), + patterns: vec![path.to_string_lossy().into_owned()], + target: Some(path.to_string_lossy().into_owned()), + command: None, + workdir: None, + workspace: dir.path().to_string_lossy().into_owned(), + reason: "approval".to_string(), + raw_input: serde_json::json!({ + "file_path": path, + "old_string": "old", + "new_string": "new", + "replace_all": false + }), + response_tx, + }; + + let content = permission_tool_content(&prompt, dir.path()); + let ToolCallContent::Diff(diff) = &content[0] else { + panic!("expected preflight diff"); + }; + assert_eq!(diff.old_text.as_deref(), Some("fn old() {}\n")); + assert_eq!(diff.new_text, "fn new() {}\n"); + } + + #[test] + fn acp_permission_apply_patch_includes_preflight_diff() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("note.txt"); + std::fs::write(&path, "before\n").unwrap(); + let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); + let prompt = crate::tools::PermissionPrompt { + tool_call_id: Some("call_patch".to_string()), + tool_id: "apply_patch".to_string(), + action: crate::tools::PermissionAction::Write, + permission: "edit".to_string(), + patterns: vec![path.to_string_lossy().into_owned()], + target: Some(path.to_string_lossy().into_owned()), + command: None, + workdir: None, + workspace: dir.path().to_string_lossy().into_owned(), + reason: "approval".to_string(), + raw_input: serde_json::json!({ + "patch": format!( + "*** Begin Patch\n*** Update File: {}\n@@\n-before\n+after\n*** End Patch\n", + path.display() + ) + }), + response_tx, + }; + + let content = permission_tool_content(&prompt, dir.path()); + let ToolCallContent::Diff(diff) = &content[0] else { + panic!("expected apply_patch preflight diff"); + }; + assert_eq!(diff.path, path); + assert_eq!(diff.old_text.as_deref(), Some("before\n")); + assert_eq!(diff.new_text, "after\n"); + } + + #[test] + fn previews_keep_workspace_diffs_but_do_not_read_protected_files() { + let workspace = tempfile::tempdir().unwrap(); + let external = tempfile::tempdir().unwrap(); + let ordinary = workspace.path().join("hello.txt"); + std::fs::write(&ordinary, "before").unwrap(); + let preview = + serde_json::to_string(&preflight_diff("hello.txt", "after", workspace.path())).unwrap(); + assert!(preview.contains("before")); + for path in [ + workspace.path().join(".env"), + external.path().join("secret.txt"), + ] { + std::fs::write(&path, "secret-never-preview").unwrap(); + let preview = serde_json::to_string(&preflight_diff( + path.to_str().unwrap(), + "after", + workspace.path(), + )) + .unwrap(); + assert!(!preview.contains("secret-never-preview")); + assert!(permission_edit_diff( + &serde_json::json!({ + "file_path": path, "old_string": "secret", "new_string": "public" + }), + workspace.path() + ) + .is_none()); + } + } + + #[cfg(unix)] + #[test] + fn previews_reject_symlink_escape_and_sensitive_alias() { + let workspace = tempfile::tempdir().unwrap(); + let external = tempfile::tempdir().unwrap(); + let secret = external.path().join("secret.txt"); + std::fs::write(&secret, "secret").unwrap(); + let link = workspace.path().join("alias.txt"); + std::os::unix::fs::symlink(secret, &link).unwrap(); + assert!(!preview_read_allowed(&link, workspace.path())); + let env = workspace.path().join(".env"); + std::fs::write(&env, "secret").unwrap(); + let alias = workspace.path().join("config.txt"); + std::os::unix::fs::symlink(env, &alias).unwrap(); + assert!(!preview_read_allowed(&alias, workspace.path())); + } +} diff --git a/src/acp/service/permissions.rs b/src/acp/service/permissions.rs index ee620ec..6f4d72a 100644 --- a/src/acp/service/permissions.rs +++ b/src/acp/service/permissions.rs @@ -1,13 +1,12 @@ -use super::notifications::absolute_tool_path; use super::notifications::tool_kind; -use super::notifications::tool_locations; +use super::permission_preview::permission_tool_content; +use super::tool_locations::tool_locations; use super::AcpSession; use crate::config::configuration::LoadedConfig; use agent_client_protocol::schema::v1::PermissionOption; use agent_client_protocol::schema::v1::PermissionOptionKind; use agent_client_protocol::schema::v1::RequestPermissionOutcome; use agent_client_protocol::schema::v1::RequestPermissionRequest; -use agent_client_protocol::schema::v1::ToolCallContent; use agent_client_protocol::schema::v1::ToolCallStatus; use agent_client_protocol::schema::v1::ToolCallUpdate; use agent_client_protocol::schema::v1::ToolCallUpdateFields; @@ -28,106 +27,6 @@ fn permission_options() -> Vec { ] } -fn permission_tool_content( - prompt: &crate::tools::PermissionPrompt, - cwd: &Path, -) -> Vec { - match prompt.tool_id.as_str() { - "write" => prompt - .raw_input - .get("file_path") - .or_else(|| prompt.raw_input.get("filePath")) - .and_then(serde_json::Value::as_str) - .zip( - prompt - .raw_input - .get("content") - .and_then(serde_json::Value::as_str), - ) - .map(|(path, new_text)| vec![preflight_diff(path, new_text, cwd)]) - .unwrap_or_default(), - "write_files" => prompt - .raw_input - .get("files") - .and_then(serde_json::Value::as_array) - .into_iter() - .flatten() - .filter_map(|file| { - let path = file.get("file_path")?.as_str()?; - let content = file.get("content")?.as_str()?; - Some(preflight_diff(path, content, cwd)) - }) - .collect(), - "edit" => permission_edit_diff(&prompt.raw_input, cwd) - .into_iter() - .collect(), - "apply_patch" - if crate::tools::patch::patch_paths_from_params(&prompt.raw_input) - .iter() - .all(|path| preview_read_allowed(&absolute_tool_path(path, cwd), cwd)) => - { - crate::tools::patch::preview_patch(&prompt.raw_input, cwd) - .unwrap_or_default() - .into_iter() - .map(|change| { - agent_client_protocol::schema::v1::Diff::new(change.path, change.new_text) - .old_text(change.old_text) - .into() - }) - .collect() - } - _ => Vec::new(), - } -} - -// A write/edit approval is not a read approval. Keep ordinary workspace previews, -// but never acquire protected file contents just to display a permission request. -fn preview_read_allowed(path: &Path, cwd: &Path) -> bool { - !crate::tools::permission::is_outside_workdir(path, cwd) - && !crate::tools::permission::is_sensitive_path(path) - && !std::fs::canonicalize(path) - .ok() - .is_some_and(|resolved| crate::tools::permission::is_sensitive_path(&resolved)) -} - -fn preflight_diff(path: &str, new_text: &str, cwd: &Path) -> ToolCallContent { - let path = absolute_tool_path(path, cwd); - let old_text = preview_read_allowed(&path, cwd) - .then(|| std::fs::read_to_string(&path).ok()) - .flatten(); - agent_client_protocol::schema::v1::Diff::new(path, new_text.to_string()) - .old_text(old_text) - .into() -} - -fn permission_edit_diff(input: &serde_json::Value, cwd: &Path) -> Option { - let path = input - .get("file_path") - .or_else(|| input.get("filePath"))? - .as_str()?; - let old_string = input.get("old_string")?.as_str()?; - let new_string = input.get("new_string")?.as_str()?; - let absolute = absolute_tool_path(path, cwd); - if !preview_read_allowed(&absolute, cwd) { - return None; - } - let old_text = std::fs::read_to_string(&absolute).ok()?; - let new_text = if input - .get("replace_all") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - { - old_text.replace(old_string, new_string) - } else { - old_text.replacen(old_string, new_string, 1) - }; - Some( - agent_client_protocol::schema::v1::Diff::new(absolute, new_text) - .old_text(old_text) - .into(), - ) -} - fn permission_tool_call_id(tool_call_id: Option<&str>) -> String { tool_call_id .map(str::to_string) @@ -219,7 +118,6 @@ mod tests { use super::super::test_support::test_session; use super::*; use agent_client_protocol::schema::v1::PermissionOptionKind; - use agent_client_protocol::schema::v1::ToolCallContent; #[test] fn acp_permission_prefers_originating_tool_call_id() { @@ -284,121 +182,4 @@ mod tests { .is_ok()); assert!(rx.try_recv().is_err()); } - - #[test] - fn acp_permission_edit_includes_preflight_diff() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("main.rs"); - std::fs::write(&path, "fn old() {}\n").unwrap(); - let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); - let prompt = crate::tools::PermissionPrompt { - tool_call_id: Some("call_edit".to_string()), - tool_id: "edit".to_string(), - action: crate::tools::PermissionAction::Write, - permission: "edit".to_string(), - patterns: vec![path.to_string_lossy().into_owned()], - target: Some(path.to_string_lossy().into_owned()), - command: None, - workdir: None, - workspace: dir.path().to_string_lossy().into_owned(), - reason: "approval".to_string(), - raw_input: serde_json::json!({ - "file_path": path, - "old_string": "old", - "new_string": "new", - "replace_all": false - }), - response_tx, - }; - - let content = permission_tool_content(&prompt, dir.path()); - let ToolCallContent::Diff(diff) = &content[0] else { - panic!("expected preflight diff"); - }; - assert_eq!(diff.old_text.as_deref(), Some("fn old() {}\n")); - assert_eq!(diff.new_text, "fn new() {}\n"); - } - - #[test] - fn acp_permission_apply_patch_includes_preflight_diff() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("note.txt"); - std::fs::write(&path, "before\n").unwrap(); - let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); - let prompt = crate::tools::PermissionPrompt { - tool_call_id: Some("call_patch".to_string()), - tool_id: "apply_patch".to_string(), - action: crate::tools::PermissionAction::Write, - permission: "edit".to_string(), - patterns: vec![path.to_string_lossy().into_owned()], - target: Some(path.to_string_lossy().into_owned()), - command: None, - workdir: None, - workspace: dir.path().to_string_lossy().into_owned(), - reason: "approval".to_string(), - raw_input: serde_json::json!({ - "patch": format!( - "*** Begin Patch\n*** Update File: {}\n@@\n-before\n+after\n*** End Patch\n", - path.display() - ) - }), - response_tx, - }; - - let content = permission_tool_content(&prompt, dir.path()); - let ToolCallContent::Diff(diff) = &content[0] else { - panic!("expected apply_patch preflight diff"); - }; - assert_eq!(diff.path, path); - assert_eq!(diff.old_text.as_deref(), Some("before\n")); - assert_eq!(diff.new_text, "after\n"); - } - - #[test] - fn previews_keep_workspace_diffs_but_do_not_read_protected_files() { - let workspace = tempfile::tempdir().unwrap(); - let external = tempfile::tempdir().unwrap(); - let ordinary = workspace.path().join("hello.txt"); - std::fs::write(&ordinary, "before").unwrap(); - let preview = - serde_json::to_string(&preflight_diff("hello.txt", "after", workspace.path())).unwrap(); - assert!(preview.contains("before")); - for path in [ - workspace.path().join(".env"), - external.path().join("secret.txt"), - ] { - std::fs::write(&path, "secret-never-preview").unwrap(); - let preview = serde_json::to_string(&preflight_diff( - path.to_str().unwrap(), - "after", - workspace.path(), - )) - .unwrap(); - assert!(!preview.contains("secret-never-preview")); - assert!(permission_edit_diff( - &serde_json::json!({ - "file_path": path, "old_string": "secret", "new_string": "public" - }), - workspace.path() - ) - .is_none()); - } - } - - #[cfg(unix)] - #[test] - fn previews_reject_symlink_escape_and_sensitive_alias() { - let workspace = tempfile::tempdir().unwrap(); - let external = tempfile::tempdir().unwrap(); - let secret = external.path().join("secret.txt"); - std::fs::write(&secret, "secret").unwrap(); - let link = workspace.path().join("alias.txt"); - std::os::unix::fs::symlink(secret, &link).unwrap(); - assert!(!preview_read_allowed(&link, workspace.path())); - let env = workspace.path().join(".env"); - std::fs::write(&env, "secret").unwrap(); - let alias = workspace.path().join("config.txt"); - std::os::unix::fs::symlink(env, &alias).unwrap(); - assert!(!preview_read_allowed(&alias, workspace.path())); - } } diff --git a/src/acp/service/prompt.rs b/src/acp/service/prompt.rs index b8763dc..d4372d4 100644 --- a/src/acp/service/prompt.rs +++ b/src/acp/service/prompt.rs @@ -1,18 +1,18 @@ +use super::attachment_storage::ManagedAttachmentGuard; use super::attachments::model_supports_audio; use super::attachments::prompt_content; -use super::attachments::ManagedAttachmentGuard; use super::cancellable; +use super::command_expansion::expand_slash_command; +use super::command_expansion::SlashExpansion; use super::commands::command_text; use super::commands::compact_command; -use super::commands::expand_slash_command; -use super::commands::SlashExpansion; -use super::config::find_selectable_model; -use super::config::model_context_window; -use super::config::resolved_reasoning; +use super::context::acp_context_tokens; use super::internal_error; -use super::notifications::send_replay_text; +use super::models::find_selectable_model; +use super::models::model_context_window; +use super::models::resolved_reasoning; use super::permissions::tool_permissions; -use super::usage::acp_context_tokens; +use super::replay::send_replay_text; use super::usage::acp_session_cost; use super::usage::send_usage; use super::AcpService; diff --git a/src/acp/service/replay.rs b/src/acp/service/replay.rs new file mode 100644 index 0000000..7821316 --- /dev/null +++ b/src/acp/service/replay.rs @@ -0,0 +1,160 @@ +use super::internal_error; +use super::notifications::send_tool_result; +use super::notifications::tool_kind; +use super::notifications::tool_title; +use super::tool_locations::tool_locations; +use agent_client_protocol::schema::v1::ContentChunk; +use agent_client_protocol::schema::v1::SessionNotification; +use agent_client_protocol::schema::v1::SessionUpdate; +use agent_client_protocol::schema::v1::ToolCall; +use agent_client_protocol::schema::v1::ToolCallStatus; +use agent_client_protocol::Client; +use agent_client_protocol::ConnectionTo; +use agent_client_protocol::Error; +use std::path::Path; + +pub(super) fn replay_messages( + connection: &ConnectionTo, + session_id: &str, + messages: &[crate::session::types::Message], + cwd: &Path, +) -> Result<(), Error> { + for message in messages { + if crate::session::compaction::is_compaction_display_item(message) { + continue; + } + let message_id = message.id.clone(); + match message.role { + crate::session::types::MessageRole::User => { + if !message.content.is_empty() { + send_replay_text( + connection, + session_id, + &message_id, + &message.content, + true, + false, + )?; + } + } + crate::session::types::MessageRole::Assistant => { + for part in &message.parts { + match part.part_type.as_str() { + "text" => { + if let Some(text) = part.text_value() { + send_replay_text( + connection, + session_id, + &message_id, + text, + false, + false, + )?; + } + } + "reasoning" => { + if let Some(text) = part.text_value() { + send_replay_text( + connection, + session_id, + &message_id, + text, + false, + true, + )?; + } + } + "tool_call" => replay_tool_call(connection, session_id, part, cwd)?, + "tool_result" => replay_tool_result(connection, session_id, part, cwd)?, + _ => {} + } + } + } + crate::session::types::MessageRole::System + | crate::session::types::MessageRole::Tool => {} + } + } + Ok(()) +} + +pub(super) fn send_replay_text( + connection: &ConnectionTo, + session_id: &str, + message_id: &str, + text: &str, + user: bool, + thought: bool, +) -> Result<(), Error> { + let chunk = ContentChunk::new(text.to_string().into()).message_id(message_id); + let update = if user { + SessionUpdate::UserMessageChunk(chunk) + } else if thought { + SessionUpdate::AgentThoughtChunk(chunk) + } else { + SessionUpdate::AgentMessageChunk(chunk) + }; + connection + .send_notification(SessionNotification::new(session_id.to_string(), update)) + .map_err(|_| internal_error()) +} + +fn replay_tool_call( + connection: &ConnectionTo, + session_id: &str, + part: &crate::session::types::MessagePart, + cwd: &Path, +) -> Result<(), Error> { + let Some(tool_call_id) = part.tool_id() else { + return Ok(()); + }; + let name = part.tool_name().unwrap_or("tool"); + let input = part + .data + .get("args") + .cloned() + .unwrap_or(serde_json::Value::Null); + let status = match part.tool_status() { + Some("completed") | Some("ok") => ToolCallStatus::Completed, + Some("error") | Some("failed") => ToolCallStatus::Failed, + Some("running") => ToolCallStatus::InProgress, + _ => ToolCallStatus::Pending, + }; + let update = SessionUpdate::ToolCall( + ToolCall::new(tool_call_id.to_string(), tool_title(name, &input)) + .kind(tool_kind(name)) + .status(status) + .locations(tool_locations(name, &input, cwd)) + .raw_input(input), + ); + connection + .send_notification(SessionNotification::new(session_id.to_string(), update)) + .map_err(|_| internal_error()) +} + +fn replay_tool_result( + connection: &ConnectionTo, + session_id: &str, + part: &crate::session::types::MessagePart, + cwd: &Path, +) -> Result<(), Error> { + let Some(tool_call_id) = part.tool_id() else { + return Ok(()); + }; + let content = part + .data + .get("content") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string(); + send_tool_result( + connection, + session_id, + crate::llm::ToolCallResult { + tool_call_id: tool_call_id.to_string(), + role: "tool".to_string(), + name: part.tool_name().unwrap_or("tool").to_string(), + content, + }, + cwd, + ) +} diff --git a/src/acp/service/sessions.rs b/src/acp/service/sessions.rs index 65542b3..5992385 100644 --- a/src/acp/service/sessions.rs +++ b/src/acp/service/sessions.rs @@ -1,12 +1,12 @@ -use super::config::merge_acp_mcp_servers; -use super::config::model_context_window; -use super::config::model_reasoning; -use super::config::resolve_model; use super::config::session_config_options; -use super::config::session_modes; use super::internal_error; -use super::notifications::replay_messages; +use super::mcp_config::merge_acp_mcp_servers; +use super::models::model_context_window; +use super::models::model_reasoning; +use super::models::resolve_model; +use super::modes::session_modes; use super::permissions::configured_tool_permissions; +use super::replay::replay_messages; use super::AcpService; use super::AcpSession; use agent_client_protocol::schema::v1::ListSessionsResponse; @@ -360,52 +360,5 @@ impl AcpService { } #[cfg(test)] -mod tests { - use super::*; - use std::path::Path; - - #[test] - fn session_info_includes_hierarchy_metadata() { - let now = std::time::SystemTime::now(); - let info = acp_session_info( - crate::session::manager::SessionInfo { - id: "child".to_string(), - parent_id: Some("parent".to_string()), - title: "Child".to_string(), - created_at: now, - updated_at: now, - message_count: 0, - workspace_id: 1, - workspace_path: "/tmp".to_string(), - workspace_name: "tmp".to_string(), - workspace_sort_order: 0, - status: crate::session::types::SessionStatus::Idle, - pinned_at: None, - archived_at: None, - }, - Some("root".to_string()), - ); - let meta = info.meta.expect("session metadata"); - assert_eq!(meta["crabcode"]["parentSessionId"], "parent"); - assert_eq!(meta["crabcode"]["rootSessionId"], "root"); - } - - #[test] - fn session_list_cursor_pages_all_sessions() { - assert_eq!( - session_page(None, 250).unwrap(), - (0, 100, Some("100".to_string())) - ); - assert_eq!( - session_page(Some("100"), 250).unwrap(), - (100, 200, Some("200".to_string())) - ); - assert_eq!(session_page(Some("200"), 250).unwrap(), (200, 250, None)); - assert!(session_page(Some("invalid"), 250).is_err()); - } - - #[test] - fn rejects_non_absolute_workspaces() { - assert!(workspace_path(Path::new("relative")).is_err()); - } -} +#[path = "tests/sessions.rs"] +mod tests; diff --git a/src/acp/service/stream.rs b/src/acp/service/stream.rs index 2a6d68e..9159a3b 100644 --- a/src/acp/service/stream.rs +++ b/src/acp/service/stream.rs @@ -8,6 +8,8 @@ use super::permissions::request_permission; use super::questions::request_questions; use super::questions::skipped_question_answers; use super::terminals::bridge_terminal_session; +use super::transcript::record_tool_call; +use super::transcript::record_tool_result; use super::usage::cumulative_cost; use super::usage::estimate_session_usage_cost; use super::usage::send_live_usage; @@ -31,86 +33,6 @@ fn acp_stop_reason(reason: Option) -> StopReason { } } -pub(super) fn record_tool_call( - assistant: &mut crate::session::types::Message, - tool_call: &crate::llm::ToolCall, -) { - let args = serde_json::from_str(&tool_call.function.arguments) - .unwrap_or_else(|_| serde_json::Value::String(tool_call.function.arguments.clone())); - let provider_executed = matches!( - tool_call.function.name.as_str(), - "x_search" | "web_search" | "file_search" - ); - - if let Some(part) = assistant - .parts - .iter_mut() - .find(|part| part.part_type == "tool_call" && part.tool_id() == Some(tool_call.id.as_str())) - { - if let Some(obj) = part.data.as_object_mut() { - obj.insert( - "name".to_string(), - serde_json::Value::String(tool_call.function.name.clone()), - ); - obj.insert("args".to_string(), args); - if provider_executed { - obj.insert( - "provider_executed".to_string(), - serde_json::Value::Bool(true), - ); - } - } - return; - } - - assistant.add_tool_call_part(tool_call.id.clone(), tool_call.function.name.clone(), args); - if provider_executed { - if let Some(obj) = assistant - .parts - .last_mut() - .and_then(|part| part.data.as_object_mut()) - { - obj.insert( - "provider_executed".to_string(), - serde_json::Value::Bool(true), - ); - } - } -} - -pub(super) fn record_tool_result( - assistant: &mut crate::session::types::Message, - result: &crate::llm::ToolCallResult, -) { - let mut data = assistant - .tool_call_part_data(&result.tool_call_id) - .cloned() - .unwrap_or_else(|| serde_json::json!({})); - data["id"] = serde_json::Value::String(result.tool_call_id.clone()); - data["name"] = serde_json::Value::String(result.name.clone()); - if matches!( - result.name.as_str(), - "x_search" | "web_search" | "file_search" - ) { - data["provider_executed"] = serde_json::Value::Bool(true); - } - - match serde_json::from_str::(&result.content) { - Ok(serde_json::Value::Object(payload)) => { - if let Some(data) = data.as_object_mut() { - for (key, value) in payload { - data.insert(key, value); - } - } - } - _ => { - data["status"] = serde_json::Value::String("ok".to_string()); - data["output_preview"] = serde_json::Value::String(result.content.clone()); - } - } - assistant.add_or_update_tool_result_part(data); -} - impl AcpService { // Keep finalization outside every cancellable await, including client requests. pub(super) async fn receive_prompt_stream( @@ -343,208 +265,5 @@ impl AcpService { } #[cfg(test)] -mod tests { - use super::super::test_support::empty_config; - use super::super::test_support::session_with_config; - use super::super::AcpService; - use super::super::AcpSessions; - use super::*; - use crate::session::manager::SessionManager; - use agent_client_protocol::schema::v1::RequestPermissionRequest; - use agent_client_protocol::schema::v1::SessionNotification; - use agent_client_protocol::schema::v1::SessionUpdate; - use agent_client_protocol::schema::v1::StopReason; - use agent_client_protocol::Client; - use std::sync::Arc; - use std::sync::Mutex; - use tokio::sync::Mutex as AsyncMutex; - - #[test] - fn maps_typed_turn_stop_reasons_to_acp() { - assert_eq!( - acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTokens)), - StopReason::MaxTokens - ); - assert_eq!( - acp_stop_reason(Some(crate::llm::TurnStopReason::Refusal)), - StopReason::Refusal - ); - assert_eq!( - acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), - StopReason::MaxTurnRequests - ); - assert_eq!( - acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), - StopReason::MaxTurnRequests - ); - assert_eq!(acp_stop_reason(None), StopReason::EndTurn); - } - - #[tokio::test(flavor = "current_thread")] - async fn cancelled_stream_persists_partial_transcript_and_interrupted_status() { - cancelled_stream_finalizes(false, false).await; - } - - #[tokio::test(flavor = "current_thread")] - async fn cancelled_permission_persists_partial_transcript_and_interrupted_status() { - cancelled_stream_finalizes(true, false).await; - } - - #[tokio::test(flavor = "current_thread")] - async fn close_reopen_waits_for_old_prompt_finalization() { - cancelled_stream_finalizes(false, true).await; - } - - async fn cancelled_stream_finalizes(pending_permission: bool, close_reopen: bool) { - use agent_client_protocol::schema::v1::RequestPermissionResponse; - use agent_client_protocol::{Agent, Responder}; - use std::time::Duration; - - tokio::task::LocalSet::new().run_until(async { - let workspace = tempfile::tempdir().unwrap(); - let mut manager = SessionManager::new() - .with_history_for_workspace(workspace.path()).unwrap(); - let session_id = manager.create_session(Some("ACP cancellation regression".into())); - let db_id = manager.get_db_id(&session_id).unwrap(); - manager.set_session_status( - &session_id, crate::session::types::SessionStatus::Streaming, None, - ).unwrap(); - let service = AcpService { - sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), - session_manager: Arc::new(Mutex::new(manager)), - client_capabilities: Arc::new(Mutex::new(Default::default())), - }; - let session = session_with_config(empty_config()); - service.sessions.lock().await.attachments.insert(session_id.clone(), session.clone()); - let cancellation = service.sessions.lock().await.begin_prompt(&session_id).unwrap(); - let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); - let (permission_tx, mut permission_rx) = tokio::sync::mpsc::unbounded_channel(); - let client = Client.builder().on_receive_request( - async move |_request: RequestPermissionRequest, - responder: Responder, _connection| { - permission_tx.send(responder).unwrap(); - Ok(()) - }, - agent_client_protocol::on_receive_request!(), - ).on_receive_notification( - async move |notification: SessionNotification, _connection| { - if matches!(notification.update, SessionUpdate::AgentMessageChunk(_)) { - let _ = observed_tx.send(()); - } - Ok(()) - }, - agent_client_protocol::on_receive_notification!(), - ); - Agent.builder().connect_with(client, async move |connection| { - let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); - sender.send(crate::llm::ChunkMessage::Text("partial answer".into())).unwrap(); - sender.send(crate::llm::ChunkMessage::Reasoning("partial reasoning".into())).unwrap(); - let (response_tx, response_rx) = tokio::sync::oneshot::channel(); - if pending_permission { - sender.send(crate::llm::ChunkMessage::PermissionRequest(crate::tools::PermissionPrompt { - tool_call_id: Some("pending-tool".into()), - tool_id: "bash".into(), - action: crate::tools::PermissionAction::Write, - permission: "bash".into(), - patterns: Vec::new(), - target: None, - command: Some("mock-command".into()), - workdir: None, - workspace: workspace.path().to_string_lossy().into_owned(), - reason: "regression test".into(), - raw_input: serde_json::json!({"command": "mock-command"}), - response_tx, - })).unwrap(); - } - let stream = service.receive_prompt_stream( - &session_id, &session, &connection, &cancellation, receiver, 0, 0, 0.0, - ); - tokio::pin!(stream); - tokio::time::timeout(Duration::from_secs(2), async { - tokio::select! { - _ = observed_rx.recv() => {}, - result = &mut stream => panic!("stream ended before cancellation: {result:?}"), - } - }).await.unwrap(); - let _pending_responder = if pending_permission { - Some(tokio::time::timeout(Duration::from_secs(2), async { - tokio::select! { - responder = permission_rx.recv() => responder.unwrap(), - result = &mut stream => panic!("stream ended before permission: {result:?}"), - } - }).await.unwrap()) - } else { - None - }; - if close_reopen { - service.close_session(&session_id).await; - assert!(cancellation.is_cancelled()); - assert!(!service.sessions.lock().await.attachments.contains_key(&session_id)); - assert!(service.available_commands(&session_id).await.is_err()); - assert!(service.set_mode(&session_id, "Build").await.is_err()); - assert!(service.set_model(&session_id, "unused").await.is_err()); - assert!(service.set_reasoning_effort(&session_id, "none").await.is_err()); - // The same guard serves load and resume. Resume must fail before - // any provider discovery, while the old stream still owns writes. - assert!(service.resume_session(session_id.clone(), workspace.path().into()).await.is_err()); - assert!(service.load_session(session_id.clone(), workspace.path().into(), connection.clone()).await.is_err()); - assert!(service.sessions.lock().await.begin_prompt(&session_id).is_err()); - } else { - cancellation.cancel(); - } - let response = tokio::time::timeout(Duration::from_secs(1), &mut stream) - .await.unwrap().unwrap(); - assert_eq!(response.stop_reason, StopReason::Cancelled); - if pending_permission { - assert!(matches!(response_rx.await.unwrap(), crate::tools::PermissionResponse::Deny)); - } - let manager = service.session_manager.lock().unwrap(); - let stored = manager.get_session_ref(&session_id).unwrap(); - assert_eq!(stored.status, crate::session::types::SessionStatus::Interrupted); - assert_eq!(stored.messages.len(), 1); - let assistant = &stored.messages[0]; - assert_eq!(assistant.content, "partial answer"); - assert!(assistant.is_complete); - assert!(assistant.was_interrupted); - drop(manager); - // A fresh DAO verifies the real database, not just the in-memory message. - let dao = crate::persistence::HistoryDAO::new_for_workspace(workspace.path()).unwrap(); - assert_eq!(dao.get_session(db_id).unwrap().unwrap().status, "interrupted"); - let messages = dao.get_messages(db_id).unwrap(); - assert_eq!(messages.len(), 1); - assert!(messages[0].parts.iter().any(|part| - part.part_type == "text" && part.data["text"] == "partial answer" - )); - drop(dao); - if close_reopen { - // Stream persistence has ended, but wrapper finalization has not: - // reopen must stay blocked through that second status write too. - assert!(service.resume_session(session_id.clone(), workspace.path().into()).await.is_err()); - service.session_manager.lock().unwrap().set_session_status( - &session_id, crate::session::types::SessionStatus::Interrupted, None, - ).unwrap(); - service.sessions.lock().await.finish_prompt(&session_id); - let mut sessions = service.sessions.lock().await; - sessions.ensure_attachable(&session_id).unwrap(); - sessions.attachments.insert(session_id.clone(), session.clone()); - let next = sessions.begin_prompt(&session_id).unwrap(); - drop(sessions); - service.session_manager.lock().unwrap().set_session_status( - &session_id, crate::session::types::SessionStatus::Streaming, None, - ).unwrap(); - assert!(!next.is_cancelled()); - assert!(service.sessions.lock().await.active_prompts.contains(&session_id)); - let dao = crate::persistence::HistoryDAO::new_for_workspace(workspace.path()).unwrap(); - assert_eq!(dao.get_session(db_id).unwrap().unwrap().status, "streaming"); - service.cancel_session(&session_id).await; - assert!(next.is_cancelled(), "the new turn retains its cancellation token"); - service.sessions.lock().await.finish_prompt(&session_id); - } else { - service.sessions.lock().await.finish_prompt(&session_id); - } - service.session_manager.lock().unwrap().try_delete_session(&session_id).unwrap(); - Ok(()) - }).await.unwrap(); - }).await; - } -} +#[path = "tests/stream.rs"] +mod tests; diff --git a/src/acp/service/terminals.rs b/src/acp/service/terminals.rs index c9c8418..a4dc43a 100644 --- a/src/acp/service/terminals.rs +++ b/src/acp/service/terminals.rs @@ -1,5 +1,5 @@ use super::cancellable; -use super::notifications::absolute_tool_path; +use super::tool_locations::absolute_tool_path; use agent_client_protocol::schema::v1::CreateTerminalRequest; use agent_client_protocol::schema::v1::KillTerminalRequest; use agent_client_protocol::schema::v1::ReleaseTerminalRequest; @@ -205,9 +205,12 @@ mod tests { use agent_client_protocol::schema::v1::CreateTerminalRequest; use agent_client_protocol::schema::v1::KillTerminalRequest; use agent_client_protocol::schema::v1::ReleaseTerminalRequest; + use agent_client_protocol::Client; use agent_client_protocol::ConnectionTo; + use std::path::Path; + use tokio_util::sync::CancellationToken; #[test] diff --git a/src/acp/service/tests/compaction.rs b/src/acp/service/tests/compaction.rs new file mode 100644 index 0000000..2ec0b4d --- /dev/null +++ b/src/acp/service/tests/compaction.rs @@ -0,0 +1,247 @@ +use super::super::compaction_summary::compacted_messages; +use super::super::context::acp_context_tokens; +use super::super::prompt::acp_request_messages; +use super::super::test_support::test_session; +use super::super::usage::acp_session_cost; +use super::super::AcpService; +use super::super::AcpSessions; +use crate::session::manager::SessionManager; +use agent_client_protocol::schema::v1::SessionNotification; +use agent_client_protocol::schema::v1::SessionUpdate; +use agent_client_protocol::Client; +use std::sync::Arc; +use std::sync::Mutex; +use tokio::sync::Mutex as AsyncMutex; + +#[test] +fn acp_compaction_discards_provider_usage_before_compaction_boundary() { + use crate::session::types::{CompactionStats, Message, MessagePart}; + let mut assistant = Message::assistant("old answer"); + assistant + .parts + .push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); + let mut summary = Message::user(format!( + "{}\nsummary", + crate::session::compaction::SUMMARY_PREFIX + )); + summary + .parts + .push(MessagePart::usage(95_000, 500, 0, 0, 0.0)); + let marker = crate::session::compaction::compaction_marker(CompactionStats { + before_tokens: 90_100, + after_tokens: 100, + before_messages: 2, + after_messages: 1, + }); + let messages = vec![Message::user("question"), assistant, summary, marker]; + + assert_eq!( + acp_context_tokens(&messages, 1_000), + crate::session::compaction::total_context_tokens(&messages) + 1_000 + ); +} + +#[test] +fn acp_compaction_honors_reserved_tokens_and_auto_setting() { + use crate::config::configuration::CompactionConfig; + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant + .parts + .push(MessagePart::usage(70_000, 0, 10_000, 0, 0.0)); + let messages = vec![Message::user("history ".repeat(2_000)), assistant]; + assert!(crate::session::compaction::select_messages_for_compaction( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + ) + .is_some()); + let config = CompactionConfig::Settings { + auto: true, + prune: false, + reserved: Some(20_000), + }; + let used = acp_context_tokens(&messages, 1_000); + assert!(crate::session::compaction::total_context_tokens(&messages) < 80_000); + assert!(crate::session::compaction::should_auto_compact( + &config, + used, + Some(100_000), + None + )); + assert!(!crate::session::compaction::should_auto_compact( + &config, + used, + Some(100_001), + None + )); + assert!(!crate::session::compaction::should_auto_compact( + &CompactionConfig::Disabled, + used, + Some(100_000), + None + )); + assert!(!crate::session::compaction::should_auto_compact( + &config, used, None, None + )); + let no_auto = CompactionConfig::Settings { + auto: false, + prune: false, + reserved: Some(20_000), + }; + assert!(!crate::session::compaction::should_auto_compact( + &no_auto, + used, + Some(100_000), + None + )); +} + +#[test] +fn acp_compaction_keeps_fresh_system_prompt_in_request() { + use crate::session::types::{CompactionStats, Message}; + let summary = Message::user(format!( + "{}\nsummary", + crate::session::compaction::SUMMARY_PREFIX + )); + let tail = Message::assistant("retained answer"); + let marker = crate::session::compaction::compaction_marker(CompactionStats { + before_tokens: 10_000, + after_tokens: 100, + before_messages: 3, + after_messages: 2, + }); + let pending = Message::user("continue"); + let history = vec![ + Message::user("archived history"), + summary.clone(), + tail.clone(), + marker, + pending.clone(), + ]; + let system = Message::system("fresh system instructions"); + let request = acp_request_messages(&history, system.clone()); + assert_eq!(request, vec![system, summary, tail, pending]); + assert_eq!( + crate::session::compaction::filter_messages_for_context(&request), + request + ); +} + +#[test] +fn acp_compaction_cost_keeps_summary_billing_without_double_counting() { + use crate::session::types::{Message, MessagePart}; + let mut assistant = Message::assistant("answer"); + assistant.cost = Some(0.25); + assistant + .parts + .push(MessagePart::usage(100, 20, 0, 0, 0.25)); + let mut summary = Message::user(format!( + "{}\nsummary", + crate::session::compaction::SUMMARY_PREFIX + )); + summary.parts.push(MessagePart::usage(100, 20, 0, 0, 0.05)); + assert!((acp_session_cost(&[assistant, summary]) - 0.30).abs() < f64::EPSILON); +} + +#[tokio::test(flavor = "current_thread")] +async fn acp_compaction_notifies_client_of_reduced_context_and_summary_cost() { + use crate::session::types::{Message, MessagePart}; + use agent_client_protocol::Agent; + tokio::task::LocalSet::new() + .run_until(async { + let workspace = tempfile::tempdir().unwrap(); + let mut manager = SessionManager::new() + .with_history_for_workspace(workspace.path()) + .unwrap(); + let session_id = manager.create_session(Some("ACP compaction regression".into())); + let mut assistant = Message::assistant("old answer"); + assistant + .parts + .push(MessagePart::usage(80_000, 200, 0, 0, 0.25)); + let messages = vec![Message::user("history ".repeat(2_000)), assistant]; + let selection = crate::session::compaction::select_messages_for_compaction( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + ) + .unwrap(); + let mut session = test_session(); + session.context_window = Some(100_000); + let summary = crate::llm::client::CompactionSummary { + text: "Retained task details".into(), + usage: crate::aisdk::chunk::TokenUsage { + input: 3_000, + output: 20, + cache_read: 0, + cache_write: 0, + }, + }; + let (mut compacted, stats) = compacted_messages( + &messages, + &selection, + &summary, + &session, + crate::session::compaction::total_context_tokens(&messages), + messages.len(), + ) + .unwrap(); + let summary_message = compacted + .iter_mut() + .find(|message| crate::session::compaction::is_compaction_summary(message)) + .unwrap(); + summary_message + .parts + .push(MessagePart::usage(0, 1, 0, 0, 0.05)); + let expected_used = + crate::session::compaction::total_context_tokens(&compacted) + 1_000; + manager + .replace_session_messages(&session_id, compacted) + .unwrap(); + let service = AcpService { + sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), + session_manager: Arc::new(Mutex::new(manager)), + client_capabilities: Arc::new(Mutex::new(Default::default())), + }; + let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); + let client = Client.builder().on_receive_notification( + async move |notification: SessionNotification, _connection| { + let _ = observed_tx.send(notification.update); + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ); + Agent + .builder() + .connect_with(client, async move |connection| { + service + .notify_compaction(&session_id, &session, &connection, stats, 1_000) + .unwrap(); + let feedback = + tokio::time::timeout(std::time::Duration::from_secs(2), observed_rx.recv()) + .await + .unwrap() + .unwrap(); + assert!(matches!(feedback, SessionUpdate::AgentMessageChunk(_))); + let update = + tokio::time::timeout(std::time::Duration::from_secs(2), observed_rx.recv()) + .await + .unwrap() + .unwrap(); + let SessionUpdate::UsageUpdate(usage) = update else { + panic!("missing usage update") + }; + assert_eq!(usage.used, expected_used as u64); + assert_eq!(usage.size, 100_000); + assert!((usage.cost.unwrap().amount - 0.30).abs() < f64::EPSILON); + service + .session_manager + .lock() + .unwrap() + .try_delete_session(&session_id) + .unwrap(); + Ok(()) + }) + .await + .unwrap(); + }) + .await; +} diff --git a/src/acp/service/tests/sessions.rs b/src/acp/service/tests/sessions.rs new file mode 100644 index 0000000..2491a7e --- /dev/null +++ b/src/acp/service/tests/sessions.rs @@ -0,0 +1,47 @@ +use super::*; +use std::path::Path; + +#[test] +fn session_info_includes_hierarchy_metadata() { + let now = std::time::SystemTime::now(); + let info = acp_session_info( + crate::session::manager::SessionInfo { + id: "child".to_string(), + parent_id: Some("parent".to_string()), + title: "Child".to_string(), + created_at: now, + updated_at: now, + message_count: 0, + workspace_id: 1, + workspace_path: "/tmp".to_string(), + workspace_name: "tmp".to_string(), + workspace_sort_order: 0, + status: crate::session::types::SessionStatus::Idle, + pinned_at: None, + archived_at: None, + }, + Some("root".to_string()), + ); + let meta = info.meta.expect("session metadata"); + assert_eq!(meta["crabcode"]["parentSessionId"], "parent"); + assert_eq!(meta["crabcode"]["rootSessionId"], "root"); +} + +#[test] +fn session_list_cursor_pages_all_sessions() { + assert_eq!( + session_page(None, 250).unwrap(), + (0, 100, Some("100".to_string())) + ); + assert_eq!( + session_page(Some("100"), 250).unwrap(), + (100, 200, Some("200".to_string())) + ); + assert_eq!(session_page(Some("200"), 250).unwrap(), (200, 250, None)); + assert!(session_page(Some("invalid"), 250).is_err()); +} + +#[test] +fn rejects_non_absolute_workspaces() { + assert!(workspace_path(Path::new("relative")).is_err()); +} diff --git a/src/acp/service/tests/stream.rs b/src/acp/service/tests/stream.rs new file mode 100644 index 0000000..337f8b3 --- /dev/null +++ b/src/acp/service/tests/stream.rs @@ -0,0 +1,203 @@ +use super::super::test_support::empty_config; +use super::super::test_support::session_with_config; +use super::super::AcpService; +use super::super::AcpSessions; +use super::*; +use crate::session::manager::SessionManager; +use agent_client_protocol::schema::v1::RequestPermissionRequest; +use agent_client_protocol::schema::v1::SessionNotification; +use agent_client_protocol::schema::v1::SessionUpdate; +use agent_client_protocol::schema::v1::StopReason; +use agent_client_protocol::Client; +use std::sync::Arc; +use std::sync::Mutex; +use tokio::sync::Mutex as AsyncMutex; + +#[test] +fn maps_typed_turn_stop_reasons_to_acp() { + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTokens)), + StopReason::MaxTokens + ); + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::Refusal)), + StopReason::Refusal + ); + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), + StopReason::MaxTurnRequests + ); + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), + StopReason::MaxTurnRequests + ); + assert_eq!(acp_stop_reason(None), StopReason::EndTurn); +} + +#[tokio::test(flavor = "current_thread")] +async fn cancelled_stream_persists_partial_transcript_and_interrupted_status() { + cancelled_stream_finalizes(false, false).await; +} + +#[tokio::test(flavor = "current_thread")] +async fn cancelled_permission_persists_partial_transcript_and_interrupted_status() { + cancelled_stream_finalizes(true, false).await; +} + +#[tokio::test(flavor = "current_thread")] +async fn close_reopen_waits_for_old_prompt_finalization() { + cancelled_stream_finalizes(false, true).await; +} + +async fn cancelled_stream_finalizes(pending_permission: bool, close_reopen: bool) { + use agent_client_protocol::schema::v1::RequestPermissionResponse; + use agent_client_protocol::{Agent, Responder}; + use std::time::Duration; + + tokio::task::LocalSet::new().run_until(async { + let workspace = tempfile::tempdir().unwrap(); + let mut manager = SessionManager::new() + .with_history_for_workspace(workspace.path()).unwrap(); + let session_id = manager.create_session(Some("ACP cancellation regression".into())); + let db_id = manager.get_db_id(&session_id).unwrap(); + manager.set_session_status( + &session_id, crate::session::types::SessionStatus::Streaming, None, + ).unwrap(); + let service = AcpService { + sessions: Arc::new(AsyncMutex::new(AcpSessions::default())), + session_manager: Arc::new(Mutex::new(manager)), + client_capabilities: Arc::new(Mutex::new(Default::default())), + }; + let session = session_with_config(empty_config()); + service.sessions.lock().await.attachments.insert(session_id.clone(), session.clone()); + let cancellation = service.sessions.lock().await.begin_prompt(&session_id).unwrap(); + let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); + let (permission_tx, mut permission_rx) = tokio::sync::mpsc::unbounded_channel(); + let client = Client.builder().on_receive_request( + async move |_request: RequestPermissionRequest, + responder: Responder, _connection| { + permission_tx.send(responder).unwrap(); + Ok(()) + }, + agent_client_protocol::on_receive_request!(), + ).on_receive_notification( + async move |notification: SessionNotification, _connection| { + if matches!(notification.update, SessionUpdate::AgentMessageChunk(_)) { + let _ = observed_tx.send(()); + } + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ); + Agent.builder().connect_with(client, async move |connection| { + let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); + sender.send(crate::llm::ChunkMessage::Text("partial answer".into())).unwrap(); + sender.send(crate::llm::ChunkMessage::Reasoning("partial reasoning".into())).unwrap(); + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + if pending_permission { + sender.send(crate::llm::ChunkMessage::PermissionRequest(crate::tools::PermissionPrompt { + tool_call_id: Some("pending-tool".into()), + tool_id: "bash".into(), + action: crate::tools::PermissionAction::Write, + permission: "bash".into(), + patterns: Vec::new(), + target: None, + command: Some("mock-command".into()), + workdir: None, + workspace: workspace.path().to_string_lossy().into_owned(), + reason: "regression test".into(), + raw_input: serde_json::json!({"command": "mock-command"}), + response_tx, + })).unwrap(); + } + let stream = service.receive_prompt_stream( + &session_id, &session, &connection, &cancellation, receiver, 0, 0, 0.0, + ); + tokio::pin!(stream); + tokio::time::timeout(Duration::from_secs(2), async { + tokio::select! { + _ = observed_rx.recv() => {}, + result = &mut stream => panic!("stream ended before cancellation: {result:?}"), + } + }).await.unwrap(); + let _pending_responder = if pending_permission { + Some(tokio::time::timeout(Duration::from_secs(2), async { + tokio::select! { + responder = permission_rx.recv() => responder.unwrap(), + result = &mut stream => panic!("stream ended before permission: {result:?}"), + } + }).await.unwrap()) + } else { + None + }; + if close_reopen { + service.close_session(&session_id).await; + assert!(cancellation.is_cancelled()); + assert!(!service.sessions.lock().await.attachments.contains_key(&session_id)); + assert!(service.available_commands(&session_id).await.is_err()); + assert!(service.set_mode(&session_id, "Build").await.is_err()); + assert!(service.set_model(&session_id, "unused").await.is_err()); + assert!(service.set_reasoning_effort(&session_id, "none").await.is_err()); + // The same guard serves load and resume. Resume must fail before + // any provider discovery, while the old stream still owns writes. + assert!(service.resume_session(session_id.clone(), workspace.path().into()).await.is_err()); + assert!(service.load_session(session_id.clone(), workspace.path().into(), connection.clone()).await.is_err()); + assert!(service.sessions.lock().await.begin_prompt(&session_id).is_err()); + } else { + cancellation.cancel(); + } + let response = tokio::time::timeout(Duration::from_secs(1), &mut stream) + .await.unwrap().unwrap(); + assert_eq!(response.stop_reason, StopReason::Cancelled); + if pending_permission { + assert!(matches!(response_rx.await.unwrap(), crate::tools::PermissionResponse::Deny)); + } + let manager = service.session_manager.lock().unwrap(); + let stored = manager.get_session_ref(&session_id).unwrap(); + assert_eq!(stored.status, crate::session::types::SessionStatus::Interrupted); + assert_eq!(stored.messages.len(), 1); + let assistant = &stored.messages[0]; + assert_eq!(assistant.content, "partial answer"); + assert!(assistant.is_complete); + assert!(assistant.was_interrupted); + drop(manager); + // A fresh DAO verifies the real database, not just the in-memory message. + let dao = crate::persistence::HistoryDAO::new_for_workspace(workspace.path()).unwrap(); + assert_eq!(dao.get_session(db_id).unwrap().unwrap().status, "interrupted"); + let messages = dao.get_messages(db_id).unwrap(); + assert_eq!(messages.len(), 1); + assert!(messages[0].parts.iter().any(|part| + part.part_type == "text" && part.data["text"] == "partial answer" + )); + drop(dao); + if close_reopen { + // Stream persistence has ended, but wrapper finalization has not: + // reopen must stay blocked through that second status write too. + assert!(service.resume_session(session_id.clone(), workspace.path().into()).await.is_err()); + service.session_manager.lock().unwrap().set_session_status( + &session_id, crate::session::types::SessionStatus::Interrupted, None, + ).unwrap(); + service.sessions.lock().await.finish_prompt(&session_id); + let mut sessions = service.sessions.lock().await; + sessions.ensure_attachable(&session_id).unwrap(); + sessions.attachments.insert(session_id.clone(), session.clone()); + let next = sessions.begin_prompt(&session_id).unwrap(); + drop(sessions); + service.session_manager.lock().unwrap().set_session_status( + &session_id, crate::session::types::SessionStatus::Streaming, None, + ).unwrap(); + assert!(!next.is_cancelled()); + assert!(service.sessions.lock().await.active_prompts.contains(&session_id)); + let dao = crate::persistence::HistoryDAO::new_for_workspace(workspace.path()).unwrap(); + assert_eq!(dao.get_session(db_id).unwrap().unwrap().status, "streaming"); + service.cancel_session(&session_id).await; + assert!(next.is_cancelled(), "the new turn retains its cancellation token"); + service.sessions.lock().await.finish_prompt(&session_id); + } else { + service.sessions.lock().await.finish_prompt(&session_id); + } + service.session_manager.lock().unwrap().try_delete_session(&session_id).unwrap(); + Ok(()) + }).await.unwrap(); + }).await; +} diff --git a/src/acp/service/tool_content.rs b/src/acp/service/tool_content.rs new file mode 100644 index 0000000..68bfb46 --- /dev/null +++ b/src/acp/service/tool_content.rs @@ -0,0 +1,242 @@ +use super::tool_locations::absolute_tool_path; +use agent_client_protocol::schema::v1::ContentBlock; +use agent_client_protocol::schema::v1::ToolCallContent; +use std::path::Path; + +fn mcp_native_tool_content(metadata: &serde_json::Value) -> Vec { + metadata + .get("mcp_result") + .and_then(|result| result.get("content")) + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter(|block| { + !matches!( + block.get("type").and_then(serde_json::Value::as_str), + Some("text") + ) + }) + .map(|block| { + serde_json::from_value::(block.clone()) + .map(ToolCallContent::from) + .unwrap_or_else(|_| { + let kind = block + .get("type") + .and_then(serde_json::Value::as_str) + .unwrap_or("unknown"); + ToolCallContent::from(format!( + "[MCP {kind} content]\n{}", + serde_json::to_string_pretty(block).unwrap_or_else(|_| block.to_string()) + )) + }) + }) + .collect() +} + +fn metadata_has_mcp_images(metadata: Option<&serde_json::Value>) -> bool { + metadata + .and_then(|metadata| metadata.get("mcp_result")) + .and_then(|result| result.get("content")) + .and_then(serde_json::Value::as_array) + .is_some_and(|content| { + content + .iter() + .any(|block| block.get("type").and_then(serde_json::Value::as_str) == Some("image")) + }) +} + +fn tool_result_text(payload: &serde_json::Value) -> String { + payload + .get("output") + .or_else(|| payload.get("output_preview")) + .or_else(|| payload.get("error")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string() +} + +pub(super) fn tool_result_content(payload: &serde_json::Value, cwd: &Path) -> Vec { + let mut content = Vec::new(); + let text = tool_result_text(payload); + if !text.is_empty() { + content.push(ToolCallContent::from(text)); + } + + if let Some(metadata) = payload.get("metadata") { + if let Some(changes) = metadata + .get("changes") + .and_then(serde_json::Value::as_array) + { + content.extend(changes.iter().filter_map(|change| tool_diff(change, cwd))); + } else if let Some(diff) = tool_diff(metadata, cwd) { + content.push(diff); + } + content.extend(mcp_native_tool_content(metadata)); + } + + if !metadata_has_mcp_images(payload.get("metadata")) { + if let Some(images) = payload.get("images").and_then(serde_json::Value::as_array) { + content.extend(images.iter().filter_map(tool_result_image)); + } + } + + content +} + +fn tool_diff(change: &serde_json::Value, cwd: &Path) -> Option { + let path = change.get("path")?.as_str()?; + let new_text = change.get("new_text")?.as_str()?; + let old_text = change.get("old_text").and_then(serde_json::Value::as_str); + Some( + agent_client_protocol::schema::v1::Diff::new( + absolute_tool_path(path, cwd), + new_text.to_string(), + ) + .old_text(old_text.map(str::to_string)) + .into(), + ) +} + +fn tool_result_image(image: &serde_json::Value) -> Option { + let data = image.get("data_url")?.as_str()?; + let media_type = image.get("media_type")?.as_str()?; + let encoded = data + .strip_prefix("data:") + .and_then(|value| value.split_once(',')) + .map(|(_, encoded)| encoded) + .unwrap_or(data); + Some(ToolCallContent::from(ContentBlock::Image( + agent_client_protocol::schema::v1::ImageContent::new(encoded, media_type), + ))) +} + +#[cfg(test)] +mod tests { + use super::super::tool_locations::tool_result_locations; + use super::*; + use agent_client_protocol::schema::v1::ContentBlock; + use agent_client_protocol::schema::v1::ToolCallContent; + + use std::path::Path; + use std::path::PathBuf; + + #[test] + fn acp_tool_result_prefers_full_output() { + let payload = serde_json::json!({ + "output": "complete tool output", + "output_preview": "short preview", + }); + + assert_eq!(tool_result_text(&payload), "complete tool output"); + } + + #[test] + fn acp_tool_result_supports_legacy_preview_payloads() { + let payload = serde_json::json!({"output_preview": "legacy output"}); + + assert_eq!(tool_result_text(&payload), "legacy output"); + } + + #[test] + fn acp_tool_result_emits_diff_location_and_image_content() { + let payload = serde_json::json!({ + "output": "updated", + "metadata": { + "path": "src/main.rs", + "line_number": 4, + "old_text": "fn old() {}", + "new_text": "fn new() {}" + }, + "images": [{ + "data_url": "data:image/png;base64,aGk=", + "media_type": "image/png" + }] + }); + let cwd = Path::new("/tmp/workspace"); + + let locations = tool_result_locations(&payload, cwd); + assert_eq!( + locations[0].path, + PathBuf::from("/tmp/workspace/src/main.rs") + ); + assert_eq!(locations[0].line, Some(4)); + + let content = tool_result_content(&payload, cwd); + let diff = content.iter().find_map(|item| match item { + ToolCallContent::Diff(diff) => Some(diff), + _ => None, + }); + let diff = diff.expect("diff content"); + assert_eq!(diff.path, PathBuf::from("/tmp/workspace/src/main.rs")); + assert_eq!(diff.old_text.as_deref(), Some("fn old() {}")); + assert_eq!(diff.new_text, "fn new() {}"); + assert!(content.iter().any(|item| matches!( + item, + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Image(image) if image.data == "aGk=" && image.mime_type == "image/png") + ))); + } + + #[test] + fn acp_tool_result_restores_native_mcp_resources() { + let metadata = serde_json::json!({ + "mcp_result": { + "content": [ + { + "type": "image", + "data": "aGk=", + "mimeType": "image/png", + "annotations": { "priority": 0.5 } + }, + { + "type": "resource_link", + "uri": "file:///tmp/readme.md", + "name": "readme", + "mimeType": "text/markdown", + "annotations": { "priority": 0.8 } + } + ] + } + }); + + let content = mcp_native_tool_content(&metadata); + assert_eq!(content.len(), 2); + assert!(matches!( + &content[0], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Image(image) + if image.data == "aGk=" && image.annotations.is_some()) + )); + assert!(matches!( + &content[1], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::ResourceLink(link) + if link.uri == "file:///tmp/readme.md") + )); + } + + #[test] + fn acp_tool_result_restores_mcp_audio_and_preserves_unknown_blocks() { + let metadata = serde_json::json!({ + "mcp_result": { + "content": [ + { "type": "audio", "data": "YXVkaW8=", "mimeType": "audio/wav" }, + { "type": "future_media", "payload": { "value": 1 } } + ] + } + }); + let content = mcp_native_tool_content(&metadata); + assert!(matches!( + &content[0], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Audio(audio) + if audio.data == "YXVkaW8=" && audio.mime_type == "audio/wav") + )); + assert!(matches!( + &content[1], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Text(text) + if text.text.contains("future_media") && text.text.contains("payload")) + )); + } +} diff --git a/src/acp/service/tool_locations.rs b/src/acp/service/tool_locations.rs new file mode 100644 index 0000000..a770eef --- /dev/null +++ b/src/acp/service/tool_locations.rs @@ -0,0 +1,116 @@ +use agent_client_protocol::schema::v1::ToolCallLocation; +use std::path::Path; +use std::path::PathBuf; + +pub(super) fn absolute_tool_path(path: &str, cwd: &Path) -> PathBuf { + let path = PathBuf::from(path); + if path.is_absolute() { + path + } else { + cwd.join(path) + } +} + +pub(super) fn tool_locations( + tool_name: &str, + input: &serde_json::Value, + cwd: &Path, +) -> Vec { + let paths = if tool_name == "write_files" { + input + .get("files") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|file| file.get("file_path").and_then(serde_json::Value::as_str)) + .map(str::to_string) + .collect() + } else if tool_name == "apply_patch" { + crate::tools::patch::patch_paths_from_params(input) + } else { + input + .get("file_path") + .or_else(|| input.get("filePath")) + .or_else(|| input.get("filepath")) + .or_else(|| input.get("path")) + .and_then(serde_json::Value::as_str) + .map(|path| vec![path.to_string()]) + .unwrap_or_default() + }; + + paths + .into_iter() + .map(|path| ToolCallLocation::new(absolute_tool_path(&path, cwd))) + .collect() +} + +pub(super) fn tool_result_locations( + payload: &serde_json::Value, + cwd: &Path, +) -> Vec { + let Some(metadata) = payload.get("metadata") else { + return Vec::new(); + }; + let line = metadata + .get("line_number") + .and_then(serde_json::Value::as_u64) + .and_then(|line| u32::try_from(line).ok()); + + if let Some(changes) = metadata + .get("changes") + .and_then(serde_json::Value::as_array) + { + return changes + .iter() + .filter_map(|change| change.get("path").and_then(serde_json::Value::as_str)) + .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd))) + .collect(); + } + + metadata + .get("path") + .and_then(serde_json::Value::as_str) + .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd)).line(line)) + .into_iter() + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::Path; + use std::path::PathBuf; + + #[test] + fn acp_tool_locations_normalize_multi_file_and_patch_paths() { + let cwd = Path::new("/tmp/workspace"); + let write_locations = tool_locations( + "write_files", + &serde_json::json!({ + "files": [ + {"file_path": "src/a.rs", "content": "a"}, + {"file_path": "/tmp/b.rs", "content": "b"} + ] + }), + cwd, + ); + assert_eq!( + write_locations[0].path, + PathBuf::from("/tmp/workspace/src/a.rs") + ); + assert_eq!(write_locations[1].path, PathBuf::from("/tmp/b.rs")); + + let patch_locations = tool_locations( + "apply_patch", + &serde_json::json!({ + "patch": "*** Begin Patch\n*** Update File: src/a.rs\n*** Add File: src/b.rs\n*** End Patch" + }), + cwd, + ); + assert_eq!(patch_locations.len(), 2); + assert_eq!( + patch_locations[1].path, + PathBuf::from("/tmp/workspace/src/b.rs") + ); + } +} diff --git a/src/acp/service/transcript.rs b/src/acp/service/transcript.rs new file mode 100644 index 0000000..3ba5148 --- /dev/null +++ b/src/acp/service/transcript.rs @@ -0,0 +1,79 @@ +pub(super) fn record_tool_call( + assistant: &mut crate::session::types::Message, + tool_call: &crate::llm::ToolCall, +) { + let args = serde_json::from_str(&tool_call.function.arguments) + .unwrap_or_else(|_| serde_json::Value::String(tool_call.function.arguments.clone())); + let provider_executed = matches!( + tool_call.function.name.as_str(), + "x_search" | "web_search" | "file_search" + ); + + if let Some(part) = assistant + .parts + .iter_mut() + .find(|part| part.part_type == "tool_call" && part.tool_id() == Some(tool_call.id.as_str())) + { + if let Some(obj) = part.data.as_object_mut() { + obj.insert( + "name".to_string(), + serde_json::Value::String(tool_call.function.name.clone()), + ); + obj.insert("args".to_string(), args); + if provider_executed { + obj.insert( + "provider_executed".to_string(), + serde_json::Value::Bool(true), + ); + } + } + return; + } + + assistant.add_tool_call_part(tool_call.id.clone(), tool_call.function.name.clone(), args); + if provider_executed { + if let Some(obj) = assistant + .parts + .last_mut() + .and_then(|part| part.data.as_object_mut()) + { + obj.insert( + "provider_executed".to_string(), + serde_json::Value::Bool(true), + ); + } + } +} + +pub(super) fn record_tool_result( + assistant: &mut crate::session::types::Message, + result: &crate::llm::ToolCallResult, +) { + let mut data = assistant + .tool_call_part_data(&result.tool_call_id) + .cloned() + .unwrap_or_else(|| serde_json::json!({})); + data["id"] = serde_json::Value::String(result.tool_call_id.clone()); + data["name"] = serde_json::Value::String(result.name.clone()); + if matches!( + result.name.as_str(), + "x_search" | "web_search" | "file_search" + ) { + data["provider_executed"] = serde_json::Value::Bool(true); + } + + match serde_json::from_str::(&result.content) { + Ok(serde_json::Value::Object(payload)) => { + if let Some(data) = data.as_object_mut() { + for (key, value) in payload { + data.insert(key, value); + } + } + } + _ => { + data["status"] = serde_json::Value::String("ok".to_string()); + data["output_preview"] = serde_json::Value::String(result.content.clone()); + } + } + assistant.add_or_update_tool_result_part(data); +} diff --git a/src/acp/service/usage.rs b/src/acp/service/usage.rs index c37f6e4..fd68465 100644 --- a/src/acp/service/usage.rs +++ b/src/acp/service/usage.rs @@ -1,3 +1,4 @@ +use super::context::estimate_acp_tokens; use super::internal_error; use super::AcpSession; use agent_client_protocol::schema::v1::Cost as AcpCost; @@ -15,64 +16,6 @@ pub(super) fn acp_session_cost(messages: &[crate::session::types::Message]) -> f .sum() } -pub(super) fn acp_context_tokens( - messages: &[crate::session::types::Message], - prompt_tokens: usize, -) -> usize { - let estimated = - crate::session::compaction::total_context_tokens(messages).saturating_add(prompt_tokens); - // Retained tail usage still describes the old context, as does the summary's billed usage. - let start = messages - .iter() - .rposition(crate::session::compaction::is_compaction_marker) - .map(|index| index + 1) - .unwrap_or(0); - for (index, message) in messages.iter().enumerate().skip(start).rev() { - if message.role != crate::session::types::MessageRole::Assistant { - continue; - } - // Each tool step bills its prompt again; cumulative usage is not a context size. - let provider_context = message - .parts - .iter() - .enumerate() - .filter_map(|(part_index, part)| { - let usage = part.recorded_usage()?; - let input = usage - .input - .saturating_add(usage.cache_read) - .saturating_add(usage.cache_write); - if input == 0 { - return None; - } - let mut tokens = - usize::try_from(input.saturating_add(usage.output)).unwrap_or(usize::MAX); - if let Some(previous_prompt) = part - .data - .get("acp_prompt_tokens") - .and_then(serde_json::Value::as_u64) - { - tokens = tokens - .saturating_sub(usize::try_from(previous_prompt).unwrap_or(usize::MAX)) - .saturating_add(prompt_tokens); - } - Some(tokens.saturating_add( - crate::session::compaction::assistant_parts_context_tokens( - &message.parts[part_index + 1..], - ), - )) - }) - .max() - .filter(|tokens| *tokens > 0); - if let Some(tokens) = provider_context { - let following = - crate::session::compaction::total_context_tokens(&messages[index + 1..]); - return estimated.max(tokens.saturating_add(following)); - } - } - estimated -} - pub(super) fn estimate_session_usage_cost( session: &AcpSession, usage: &crate::aisdk::chunk::TokenUsage, @@ -134,35 +77,6 @@ pub(super) struct AcpLiveUsage { pub(super) cost: f64, } -impl AcpLiveUsage { - pub(super) fn observe_provider_usage( - &mut self, - assistant: &crate::session::types::Message, - usage: crate::aisdk::chunk::TokenUsage, - ) { - self.cumulative = self.cumulative.saturating_add(usage); - let provider_input = usage - .input - .saturating_add(usage.cache_read) - .saturating_add(usage.cache_write); - let observed_assistant = live_assistant_context_tokens(assistant); - self.provider_context_floor = self - .provider_context_floor - .max(usize::try_from(provider_input).unwrap_or(usize::MAX)) - .max(observed_assistant); - } - - pub(super) fn context_tokens( - &self, - base_context_tokens: usize, - assistant: &crate::session::types::Message, - ) -> usize { - base_context_tokens - .saturating_add(live_assistant_context_tokens(assistant)) - .max(self.provider_context_floor) - } -} - fn live_assistant_context_tokens(assistant: &crate::session::types::Message) -> usize { let persisted = crate::session::compaction::message_context_tokens(assistant); let reasoning = assistant @@ -173,10 +87,6 @@ fn live_assistant_context_tokens(assistant: &crate::session::types::Message) -> persisted.saturating_add(reasoning) } -fn estimate_acp_tokens(content: &str) -> usize { - content.chars().count().saturating_add(3) / 4 -} - pub(super) fn cumulative_cost(base_cost: f64, turn_cost: f64) -> Option { let cost = base_cost + turn_cost; (cost > 0.0).then_some(cost) @@ -210,14 +120,45 @@ fn usage_update( .meta(meta) } +impl AcpLiveUsage { + pub(super) fn observe_provider_usage( + &mut self, + assistant: &crate::session::types::Message, + usage: crate::aisdk::chunk::TokenUsage, + ) { + self.cumulative = self.cumulative.saturating_add(usage); + let provider_input = usage + .input + .saturating_add(usage.cache_read) + .saturating_add(usage.cache_write); + let observed_assistant = live_assistant_context_tokens(assistant); + self.provider_context_floor = self + .provider_context_floor + .max(usize::try_from(provider_input).unwrap_or(usize::MAX)) + .max(observed_assistant); + } + + pub(super) fn context_tokens( + &self, + base_context_tokens: usize, + assistant: &crate::session::types::Message, + ) -> usize { + base_context_tokens + .saturating_add(live_assistant_context_tokens(assistant)) + .max(self.provider_context_floor) + } +} + #[cfg(test)] mod tests { - use super::super::stream::record_tool_call; - use super::super::stream::record_tool_result; + use super::super::test_support::session_with_config; use super::super::test_support::test_session; + use super::super::transcript::record_tool_call; + use super::super::transcript::record_tool_result; use super::*; use crate::config::configuration::LoadedConfig; + use std::path::PathBuf; #[test] @@ -320,136 +261,6 @@ mod tests { assert_eq!(live.cumulative.cache_read, 60_000); } - #[test] - fn acp_compaction_counts_provider_context_and_cached_input() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant - .parts - .push(MessagePart::usage(60_000, 200, 20_000, 5_000, 0.0)); - let messages = vec![Message::user("question"), assistant]; - - assert_eq!(acp_context_tokens(&messages, 1_000), 85_200); - } - - #[test] - fn acp_compaction_does_not_sum_context_across_tool_requests() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant - .parts - .push(MessagePart::usage(60_000, 200, 20_000, 0, 0.0)); - assistant - .parts - .push(MessagePart::usage(10_000, 100, 40_000, 0, 0.0)); - - assert_eq!( - acp_context_tokens(&[Message::user("question"), assistant], 0), - 80_200 - ); - } - - #[test] - fn acp_compaction_uses_latest_turn_usage_and_counts_pending_history() { - use crate::session::types::{Message, MessagePart}; - let mut old = Message::assistant("old answer"); - old.parts.push(MessagePart::usage(90_000, 100, 0, 0, 0.0)); - let mut latest = Message::assistant("new answer"); - latest - .parts - .push(MessagePart::usage(50_000, 200, 0, 0, 0.0)); - let pending = Message::user("continue"); - let pending_tokens = crate::session::compaction::message_context_tokens(&pending); - let messages = vec![Message::user("question"), old, latest, pending]; - - assert_eq!(acp_context_tokens(&messages, 0), 50_200 + pending_tokens); - } - - #[test] - fn acp_compaction_counts_prompt_and_pending_tokens_without_provider_usage() { - use crate::config::configuration::CompactionConfig; - use crate::session::types::Message; - let messages = vec![Message::user("history ".repeat(2_000))]; - let history_tokens = crate::session::compaction::total_context_tokens(&messages); - let config = CompactionConfig::Settings { - auto: true, - prune: false, - reserved: Some(1_000), - }; - let context_window = Some(history_tokens as u32 + 3_000); - assert!(!crate::session::compaction::should_auto_compact( - &config, - acp_context_tokens(&messages, 1_500), - context_window, - None, - )); - assert!(crate::session::compaction::should_auto_compact( - &config, - acp_context_tokens(&messages, 1_500) + 500, - context_window, - None, - )); - } - - #[test] - fn acp_compaction_replaces_previous_prompt_estimate_in_provider_context() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - let mut usage = MessagePart::usage(80_000, 200, 0, 0, 0.0); - usage.data["acp_prompt_tokens"] = serde_json::json!(5_000); - assistant.parts.push(usage); - let messages = vec![Message::user("question"), assistant]; - assert_eq!(acp_context_tokens(&messages, 25_000), 100_200); - assert_eq!(acp_context_tokens(&messages, 1_000), 76_200); - } - - #[test] - fn acp_compaction_keeps_provider_context_when_new_input_usage_is_missing() { - use crate::session::types::{Message, MessagePart}; - for output in [0, 100] { - let mut previous = Message::assistant("previous answer"); - previous - .parts - .push(MessagePart::usage(85_000, 200, 0, 0, 0.0)); - let mut latest = Message::assistant("new answer"); - let mut usage = MessagePart::usage(0, output, 0, 0, 0.0); - usage.data["acp_prompt_tokens"] = serde_json::json!(1_000); - latest.parts.push(usage); - let following_tokens = crate::session::compaction::message_context_tokens(&latest); - assert_eq!( - acp_context_tokens(&[Message::user("question"), previous, latest], 2_000), - 85_200 + following_tokens - ); - } - } - - #[test] - fn acp_compaction_counts_tool_output_after_provider_usage() { - use crate::session::types::{Message, MessagePart}; - let mut assistant = Message::assistant("answer"); - assistant - .parts - .push(MessagePart::usage(80_000, 200, 0, 0, 0.0)); - record_tool_result( - &mut assistant, - &crate::llm::ToolCallResult { - tool_call_id: "call_1".into(), - role: "tool".into(), - name: "read".into(), - content: serde_json::json!({"output_preview": "late tool output ".repeat(1_000)}) - .to_string(), - }, - ); - let mut following = Message::assistant(""); - following.parts = assistant.parts[2..].to_vec(); - let extra_tokens = crate::session::compaction::message_context_tokens(&following); - assert!(extra_tokens > 1_000); - assert_eq!( - acp_context_tokens(&[Message::user("question"), assistant], 0), - 80_200 + extra_tokens - ); - } - #[test] fn unknown_model_usage_cost_is_zero() { let mut session = session_with_config(LoadedConfig { From bbea75211ff9a36408ed3752244c084d3a5b72bf Mon Sep 17 00:00:00 2001 From: Yanuar Date: Tue, 13 Oct 2026 12:20:44 +0700 Subject: [PATCH 5/5] test(acp): isolate mock compaction from runtime discovery --- tests/acp_stdio.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/acp_stdio.rs b/tests/acp_stdio.rs index 26a17a2..e357775 100644 --- a/tests/acp_stdio.rs +++ b/tests/acp_stdio.rs @@ -663,6 +663,12 @@ fn automatically_compacts_provider_usage_over_stdio() { std::fs::write(workspace.path().join("AGENTS.md"), system_marker).unwrap(); // Request-time provider discovery also reads the process working directory. let mut process = AcpProcess::spawn_with_working_directory(workspace.path(), workspace.path()); + // Empty auth storage enables runtime-provider discovery; this fixture uses only mock. + std::fs::write( + process._state.path().join("crabcode/auth.json"), + serde_json::json!({"mock": {"type": "api", "key": "dummy"}}).to_string(), + ) + .expect("mock auth fixture"); assert!(initialize(&mut process).get("error").is_none()); process.send(serde_json::json!({ "jsonrpc": "2.0", "id": 2, "method": "session/new",