diff --git a/AGENTS.md b/AGENTS.md index ef146b47..aa44c88b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,6 +1,6 @@ -# CLAUDE.md +# Repository Guidance -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. +This file provides guidance to AI agents when working with code in this repository. ## Essential Commands diff --git a/README.md b/README.md index 9ecbe79c..b4c95fd0 100644 --- a/README.md +++ b/README.md @@ -196,6 +196,21 @@ The code-assistant uses two JSON configuration files to manage LLM providers and **Full Examples**: See [`providers.example.json`](providers.example.json) and [`models.example.json`](models.example.json) for complete configuration examples with all supported providers (Anthropic, OpenAI, Ollama, SAP AI Core, Vertex AI, Groq, Cerebras, MistralAI, OpenRouter). +### Tool Configuration + +Some tools require external API keys to function. Configure these in `~/.config/code-assistant/tools.json`: + +```json +{ + "perplexity_api_key": "${PERPLEXITY_API_KEY}" +} +``` + +**Available Tool Settings**: +- `perplexity_api_key` - Enables the `perplexity_ask` tool for AI-powered web search + +Tools without their required configuration will not be available to the assistant. + **List Available Models**: ```bash # See all configured models @@ -217,7 +232,6 @@ Configure in Claude Desktop settings (**Developer** tab → **Edit Config**): "command": "/path/to/code-assistant/target/release/code-assistant", "args": ["server"], "env": { - "PERPLEXITY_API_KEY": "pplx-...", // Optional, enables perplexity_ask tool "SHELL": "/bin/zsh" // Your login shell } } diff --git a/crates/code_assistant/src/acp/agent.rs b/crates/code_assistant/src/acp/agent.rs index 1bde9e31..290c5288 100644 --- a/crates/code_assistant/src/acp/agent.rs +++ b/crates/code_assistant/src/acp/agent.rs @@ -593,11 +593,20 @@ impl acp::Agent for ACPAgentImpl { let use_acp_fs = filesystem_supported && client_connection.is_some(); + // Get the ACP root path for this session (the working directory opened in Zed) + let acp_root = { + let manager = session_manager.lock().await; + manager + .get_session(&arguments.session_id.0) + .and_then(|session| session.session.config.init_path.clone()) + }; + // Create project manager and command executor let project_manager: Box = if use_acp_fs { Box::new(AcpProjectManager::new( DefaultProjectManager::new(), arguments.session_id.clone(), + acp_root, )) } else { Box::new(DefaultProjectManager::new()) diff --git a/crates/code_assistant/src/acp/explorer.rs b/crates/code_assistant/src/acp/explorer.rs index 6e82a30f..530e34aa 100644 --- a/crates/code_assistant/src/acp/explorer.rs +++ b/crates/code_assistant/src/acp/explorer.rs @@ -330,11 +330,23 @@ impl CodeExplorer for AcpCodeExplorer { pub struct AcpProjectManager { inner: DefaultProjectManager, session_id: acp::SessionId, + /// The root directory of the ACP session (i.e., the project opened in Zed). + /// Only projects with a path matching this root will use the ACP explorer. + acp_root: Option, } impl AcpProjectManager { - pub fn new(inner: DefaultProjectManager, session_id: acp::SessionId) -> Self { - Self { inner, session_id } + pub fn new( + inner: DefaultProjectManager, + session_id: acp::SessionId, + acp_root: Option, + ) -> Self { + let acp_root = acp_root.and_then(|p| p.canonicalize().ok()); + Self { + inner, + session_id, + acp_root, + } } fn maybe_project(&self, name: &str) -> Result { @@ -342,6 +354,15 @@ impl AcpProjectManager { .get_project(name)? .ok_or_else(|| anyhow!("Project not found: {name}")) } + + /// Returns true if the given project path matches the ACP session root. + fn is_acp_project(&self, project_path: &Path) -> bool { + let Some(acp_root) = &self.acp_root else { + return false; + }; + let canonical = project_path.canonicalize().ok(); + canonical.as_ref() == Some(acp_root) + } } impl ProjectManager for AcpProjectManager { @@ -359,9 +380,22 @@ impl ProjectManager for AcpProjectManager { fn get_explorer_for_project(&self, name: &str) -> Result> { let project = self.maybe_project(name)?; - Ok(Box::new(AcpCodeExplorer::new( - project.path, - self.session_id.clone(), - ))) + + // Use ACP explorer only for the project that matches the ACP session root. + // Other projects (e.g., referenced via settings) use the standard local explorer + // because Zed's ACP filesystem access is restricted to the current project. + if self.is_acp_project(&project.path) { + Ok(Box::new(AcpCodeExplorer::new( + project.path, + self.session_id.clone(), + ))) + } else { + tracing::debug!( + "Using local explorer for project '{}' at {} (outside ACP root)", + name, + project.path.display() + ); + Ok(Box::new(Explorer::new(project.path))) + } } } diff --git a/crates/code_assistant/src/acp/types.rs b/crates/code_assistant/src/acp/types.rs index bb9b25e2..1b0e888e 100644 --- a/crates/code_assistant/src/acp/types.rs +++ b/crates/code_assistant/src/acp/types.rs @@ -31,6 +31,7 @@ pub fn fragment_to_content_block(fragment: &DisplayFragment) -> acp::ContentBloc meta: None, }) } + // Tool-related fragments are not converted to content blocks // They are handled separately as ToolCall updates DisplayFragment::ToolName { .. } @@ -40,7 +41,8 @@ pub fn fragment_to_content_block(fragment: &DisplayFragment) -> acp::ContentBloc | DisplayFragment::ToolTerminal { .. } | DisplayFragment::ReasoningSummaryStart | DisplayFragment::ReasoningSummaryDelta(_) - | DisplayFragment::ReasoningComplete => { + | DisplayFragment::ReasoningComplete + | DisplayFragment::HiddenToolCompleted => { // These should not be converted to content blocks // Return empty text as placeholder acp::ContentBlock::Text(acp::TextContent { diff --git a/crates/code_assistant/src/acp/ui.rs b/crates/code_assistant/src/acp/ui.rs index 78f2a400..bc0f80b5 100644 --- a/crates/code_assistant/src/acp/ui.rs +++ b/crates/code_assistant/src/acp/ui.rs @@ -11,6 +11,14 @@ use serde_json::{Map as JsonMap, Value as JsonValue}; use crate::acp::types::{fragment_to_content_block, map_tool_kind, map_tool_status}; use crate::ui::{DisplayFragment, UIError, UiEvent, UserInterface}; +/// Tracks the last type of content for paragraph breaks after hidden tools +#[derive(Debug, Clone, Copy, PartialEq)] +enum LastContentType { + None, + Text, + Thinking, +} + /// UserInterface implementation that sends session/update notifications via ACP pub struct ACPUserUI { session_id: acp::SessionId, @@ -21,6 +29,10 @@ pub struct ACPUserUI { // Track if we should continue streaming (atomic for lock-free access from sync callbacks) should_continue: Arc, last_error: Arc>>, + // Track last content type for paragraph breaks after hidden tools + last_content_type: Arc>, + // Flag indicating a hidden tool completed and we may need a paragraph break + needs_paragraph_break_after_hidden_tool: Arc>, } #[derive(Default, Clone)] @@ -427,6 +439,8 @@ impl ACPUserUI { base_path, should_continue: Arc::new(AtomicBool::new(true)), last_error: Arc::new(Mutex::new(None)), + last_content_type: Arc::new(Mutex::new(LastContentType::None)), + needs_paragraph_break_after_hidden_tool: Arc::new(Mutex::new(false)), } } @@ -520,6 +534,38 @@ impl ACPUserUI { } } + /// Check if we need a paragraph break after a hidden tool and emit it if so + fn maybe_emit_paragraph_break(&self, current_type: LastContentType) -> Result<(), UIError> { + let mut needs_break = self.needs_paragraph_break_after_hidden_tool.lock().unwrap(); + if !*needs_break { + return Ok(()); + } + + // Reset the flag + *needs_break = false; + + // Check if the content type matches the last one + let last_type = *self.last_content_type.lock().unwrap(); + if last_type == current_type { + // Same type as before the hidden tool - emit paragraph break + let content = acp::ContentBlock::Text(acp::TextContent { + annotations: None, + text: "\n\n".to_string(), + meta: None, + }); + let chunk = Self::content_chunk(content); + + // Use the appropriate update type based on content type + let update = match current_type { + LastContentType::Thinking => acp::SessionUpdate::AgentThoughtChunk(chunk), + _ => acp::SessionUpdate::AgentMessageChunk(chunk), + }; + self.queue_session_update(update); + } + + Ok(()) + } + pub fn tool_call_update(&self, tool_id: &str) -> Option { let base_path = self.base_path.as_deref(); let tool_calls = self.tool_calls.lock().unwrap(); @@ -709,7 +755,8 @@ impl UserInterface for ACPUserUI { | UiEvent::ClearError | UiEvent::UpdateCurrentModel { .. } | UiEvent::UpdateSandboxPolicy { .. } - | UiEvent::CancelSubAgent { .. } => { + | UiEvent::CancelSubAgent { .. } + | UiEvent::HiddenToolCompleted => { // These are UI management events, not relevant for ACP } UiEvent::DisplayError { message } => { @@ -724,7 +771,22 @@ impl UserInterface for ACPUserUI { fn display_fragment(&self, fragment: &DisplayFragment) -> Result<(), UIError> { match fragment { - DisplayFragment::PlainText(_) | DisplayFragment::Image { .. } => { + DisplayFragment::PlainText(text) => { + // Check if we need a paragraph break after a hidden tool + self.maybe_emit_paragraph_break(LastContentType::Text)?; + + // Track content type for future hidden tool events + *self.last_content_type.lock().unwrap() = LastContentType::Text; + + let content = acp::ContentBlock::Text(acp::TextContent { + annotations: None, + text: text.clone(), + meta: None, + }); + let chunk = Self::content_chunk(content); + self.queue_session_update(acp::SessionUpdate::AgentMessageChunk(chunk)); + } + DisplayFragment::Image { .. } => { let content = fragment_to_content_block(fragment); let chunk = Self::content_chunk(content); self.queue_session_update(acp::SessionUpdate::AgentMessageChunk(chunk)); @@ -734,8 +796,18 @@ impl UserInterface for ACPUserUI { let chunk = Self::content_chunk(content); self.queue_session_update(acp::SessionUpdate::AgentMessageChunk(chunk)); } - DisplayFragment::ThinkingText(_) => { - let content = fragment_to_content_block(fragment); + DisplayFragment::ThinkingText(text) => { + // Check if we need a paragraph break after a hidden tool + self.maybe_emit_paragraph_break(LastContentType::Thinking)?; + + // Track content type for future hidden tool events + *self.last_content_type.lock().unwrap() = LastContentType::Thinking; + + let content = acp::ContentBlock::Text(acp::TextContent { + annotations: None, + text: text.clone(), + meta: None, + }); let chunk = Self::content_chunk(content); self.queue_session_update(acp::SessionUpdate::AgentThoughtChunk(chunk)); } @@ -837,8 +909,13 @@ impl UserInterface for ACPUserUI { self.queue_session_update(acp::SessionUpdate::ToolCallUpdate(tool_call_update)); } + DisplayFragment::ReasoningSummaryStart | DisplayFragment::ReasoningComplete => { - // No ACP representation needed yet + // No ACP representation needed + } + DisplayFragment::HiddenToolCompleted => { + // Mark that a hidden tool completed - paragraph break may be needed before next text + *self.needs_paragraph_break_after_hidden_tool.lock().unwrap() = true; } DisplayFragment::ReasoningSummaryDelta(delta) => { // Reasoning summaries are emitted as AgentThoughtChunk, same as ThinkingText diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index 3a84c9c1..38ad7e26 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -17,7 +17,6 @@ use llm::{ }; use std::collections::HashMap; use std::fs; -use std::path::PathBuf; use std::sync::{Arc, Mutex}; use std::time::SystemTime; use tracing::{debug, trace, warn}; @@ -1024,9 +1023,10 @@ impl Agent { /// Attempt to read AGENTS.md or CLAUDE.md from the initial project root. /// Prefers AGENTS.md when both exist. Returns (file_name, content) on success. fn read_repository_guidance(&self) -> Option<(String, String)> { - // Determine search root - let root_path = if !self.session_config.initial_project.is_empty() { - PathBuf::from(&self.session_config.initial_project) + // Determine search root from init_path (the actual directory path), + // not initial_project (which is just the project name) + let root_path = if let Some(path) = &self.session_config.init_path { + path.clone() } else { std::env::current_dir().ok()? }; diff --git a/crates/code_assistant/src/persistence.rs b/crates/code_assistant/src/persistence.rs index 021693a1..b807e75d 100644 --- a/crates/code_assistant/src/persistence.rs +++ b/crates/code_assistant/src/persistence.rs @@ -353,6 +353,75 @@ impl FileSessionPersistence { Ok(()) } + /// Delete all empty sessions (sessions with no messages). + /// Returns the number of deleted sessions. + /// + /// This method uses metadata to identify potentially empty sessions first, + /// avoiding the need to load all session files. For safety, it verifies + /// each candidate session is actually empty before deleting. + pub fn delete_empty_sessions(&mut self) -> Result { + let sessions_dir = self.root_dir.join("sessions"); + if !sessions_dir.exists() { + return Ok(0); + } + + // Use metadata to find candidate empty sessions (message_count == 0) + let metadata_list = self.list_chat_sessions()?; + let candidate_ids: Vec = metadata_list + .iter() + .filter(|m| m.message_count == 0) + .map(|m| m.id.clone()) + .collect(); + + if candidate_ids.is_empty() { + return Ok(0); + } + + let mut deleted_count = 0; + + // Verify and delete each candidate + for session_id in candidate_ids { + // Safety check: load the session and verify it's actually empty + match self.load_chat_session(&session_id) { + Ok(Some(session)) if session.messages.is_empty() => { + if let Err(e) = self.delete_chat_session(&session_id) { + warn!("Failed to delete empty session {}: {}", session_id, e); + } else { + info!("Deleted empty session: {}", session_id); + deleted_count += 1; + } + } + Ok(Some(_)) => { + // Metadata was out of sync, session has messages - skip + debug!( + "Session {} has messages despite metadata saying 0, skipping", + session_id + ); + } + Ok(None) => { + // Session file doesn't exist, clean up metadata + debug!( + "Session {} file not found, cleaning up metadata", + session_id + ); + let _ = self.delete_chat_session(&session_id); + } + Err(e) => { + warn!( + "Failed to load session {} for verification: {}", + session_id, e + ); + } + } + } + + if deleted_count > 0 { + info!("Cleaned up {} empty session(s)", deleted_count); + } + + Ok(deleted_count) + } + /// Rebuild metadata from existing session files (used when metadata file is corrupted) fn rebuild_metadata_from_sessions(&self) -> Result> { let mut metadata_list = Vec::new(); @@ -681,6 +750,7 @@ mod tests { use super::*; use crate::session::SessionConfig; use crate::types::{PlanItem, PlanItemPriority, PlanItemStatus}; + use tempfile::tempdir; #[test] fn chat_session_plan_roundtrip() { @@ -711,4 +781,85 @@ mod tests { let meta = restored.plan.meta.expect("plan meta should exist"); assert_eq!(meta["source"], "unit-test"); } + + #[test] + fn delete_empty_sessions_removes_only_empty() { + // Create a temporary directory for this test + let temp_dir = tempdir().expect("failed to create temp dir"); + + // Create a persistence instance using the temp directory + let mut persistence = FileSessionPersistence { + root_dir: temp_dir.path().to_path_buf(), + }; + + // Create an empty session + let empty_session = ChatSession::new_empty( + "empty_session".to_string(), + "Empty Session".to_string(), + SessionConfig::default(), + None, + ); + persistence + .save_chat_session(&empty_session) + .expect("save empty session"); + + // Create a session with messages + let mut non_empty_session = ChatSession::new_empty( + "non_empty_session".to_string(), + "Non-Empty Session".to_string(), + SessionConfig::default(), + None, + ); + non_empty_session.messages.push(Message::new_user("Hello")); + persistence + .save_chat_session(&non_empty_session) + .expect("save non-empty session"); + + // Verify both sessions exist + let sessions = persistence.list_chat_sessions().expect("list sessions"); + assert_eq!(sessions.len(), 2); + + // Delete empty sessions + let deleted_count = persistence + .delete_empty_sessions() + .expect("delete empty sessions"); + + // Should have deleted exactly one session + assert_eq!(deleted_count, 1); + + // Verify only the non-empty session remains + let remaining_sessions = persistence.list_chat_sessions().expect("list sessions"); + assert_eq!(remaining_sessions.len(), 1); + assert_eq!(remaining_sessions[0].id, "non_empty_session"); + + // Verify the empty session file is gone + assert!(persistence + .load_chat_session("empty_session") + .expect("load") + .is_none()); + + // Verify the non-empty session still exists + assert!(persistence + .load_chat_session("non_empty_session") + .expect("load") + .is_some()); + } + + #[test] + fn delete_empty_sessions_handles_no_sessions() { + // Create a temporary directory for this test + let temp_dir = tempdir().expect("failed to create temp dir"); + + // Create a persistence instance using the temp directory + let mut persistence = FileSessionPersistence { + root_dir: temp_dir.path().to_path_buf(), + }; + + // Delete empty sessions when there are no sessions + let deleted_count = persistence + .delete_empty_sessions() + .expect("delete empty sessions"); + + assert_eq!(deleted_count, 0); + } } diff --git a/crates/code_assistant/src/session/manager.rs b/crates/code_assistant/src/session/manager.rs index d079313c..b90b423c 100644 --- a/crates/code_assistant/src/session/manager.rs +++ b/crates/code_assistant/src/session/manager.rs @@ -18,7 +18,7 @@ use crate::ui::UserInterface; use command_executor::{CommandExecutor, SandboxedCommandExecutor}; use llm::LLMProvider; use sandbox::SandboxPolicy; -use tracing::{debug, error}; +use tracing::{debug, error, info, warn}; /// The main SessionManager that manages multiple active sessions with on-demand agents pub struct SessionManager { @@ -40,12 +40,27 @@ pub struct SessionManager { } impl SessionManager { - /// Create a new SessionManager + /// Create a new SessionManager. + /// + /// On creation, this will clean up any empty sessions from previous runs. + /// This handles the case where a client (e.g., Zed) starts code-assistant + /// and creates a session, but the user never sends a message before closing. pub fn new( - persistence: FileSessionPersistence, + mut persistence: FileSessionPersistence, session_config_template: SessionConfig, default_model_name: String, ) -> Self { + // Clean up empty sessions from previous runs at startup + match persistence.delete_empty_sessions() { + Ok(count) if count > 0 => { + info!("Cleaned up {} empty session(s) from previous runs", count); + } + Ok(_) => {} + Err(e) => { + warn!("Failed to clean up empty sessions at startup: {}", e); + } + } + Self { persistence, active_sessions: HashMap::new(), diff --git a/crates/code_assistant/src/tests/mocks.rs b/crates/code_assistant/src/tests/mocks.rs index 0eaa8c08..1d3fe6e9 100644 --- a/crates/code_assistant/src/tests/mocks.rs +++ b/crates/code_assistant/src/tests/mocks.rs @@ -298,12 +298,16 @@ impl UserInterface for MockUI { .unwrap() .push("\n• Reasoning Complete".to_string()); } + crate::ui::DisplayFragment::CompactionDivider { summary } => { self.streaming .lock() .unwrap() .push(format!("[compaction] {summary}")); } + crate::ui::DisplayFragment::HiddenToolCompleted => { + // Hidden tool completed - UI handles paragraph breaks + } } Ok(()) } diff --git a/crates/code_assistant/src/tests/sub_agent_tests.rs b/crates/code_assistant/src/tests/sub_agent_tests.rs index 89629917..7f4c7aff 100644 --- a/crates/code_assistant/src/tests/sub_agent_tests.rs +++ b/crates/code_assistant/src/tests/sub_agent_tests.rs @@ -298,7 +298,8 @@ fn test_tool_scope_for_sub_agent() { assert!(available.contains(&"glob_files".to_string())); assert!(available.contains(&"web_fetch".to_string())); assert!(available.contains(&"web_search".to_string())); - assert!(available.contains(&"perplexity_ask".to_string())); + // Note: perplexity_ask is only available if PERPLEXITY_API_KEY is configured in tools.json + // So we don't assert its presence here - it depends on runtime configuration // Check write tools are NOT available in SubAgentReadOnly scope assert!(!available.contains(&"write_file".to_string())); diff --git a/crates/code_assistant/src/tools/core/config.rs b/crates/code_assistant/src/tools/core/config.rs new file mode 100644 index 00000000..8bb74db5 --- /dev/null +++ b/crates/code_assistant/src/tools/core/config.rs @@ -0,0 +1,141 @@ +//! Configuration for tools that require external API keys or settings. + +use anyhow::{Context, Result}; +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use std::sync::OnceLock; + +/// Configuration for tools that require external services. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ToolsConfig { + /// API key for Perplexity service (enables perplexity_ask tool) + #[serde(default)] + pub perplexity_api_key: Option, +} + +impl ToolsConfig { + /// Get the global singleton instance of the tools configuration. + /// Returns a default (empty) config if no configuration file exists. + pub fn global() -> &'static Self { + static INSTANCE: OnceLock = OnceLock::new(); + INSTANCE.get_or_init(|| Self::load().unwrap_or_default()) + } + + /// Load the tools configuration from disk. + /// Returns Ok with default config if the file doesn't exist. + pub fn load() -> Result { + let config_path = Self::config_path()?; + + if !config_path.exists() { + return Ok(Self::default()); + } + + let content = std::fs::read_to_string(&config_path) + .with_context(|| format!("Failed to read tools config: {}", config_path.display()))?; + + let config: ToolsConfig = serde_json::from_str(&content) + .with_context(|| format!("Failed to parse tools config: {}", config_path.display()))?; + + // Substitute environment variables in API keys + Self::substitute_env_vars(config) + } + + /// Get the path to the tools configuration file. + pub fn config_path() -> Result { + // Use the same config directory as other code-assistant configs + let config_dir = Self::config_directory()?; + Ok(config_dir.join("tools.json")) + } + + /// Get the configuration directory. + fn config_directory() -> Result { + // Check for custom config directory first + if let Ok(custom_dir) = std::env::var("CODE_ASSISTANT_CONFIG_DIR") { + return Ok(PathBuf::from(custom_dir)); + } + + // Check XDG_CONFIG_HOME + if let Ok(xdg_config) = std::env::var("XDG_CONFIG_HOME") { + return Ok(PathBuf::from(xdg_config).join("code-assistant")); + } + + // Fall back to ~/.config/code-assistant + let home = dirs::home_dir() + .ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?; + Ok(home.join(".config").join("code-assistant")) + } + + /// Substitute environment variables in configuration values. + /// Supports ${VAR_NAME} syntax. + fn substitute_env_vars(mut config: Self) -> Result { + if let Some(ref key) = config.perplexity_api_key { + config.perplexity_api_key = Some(Self::substitute_env_var_in_string(key)?); + } + Ok(config) + } + + /// Substitute environment variables in a string. + fn substitute_env_var_in_string(input: &str) -> Result { + let mut result = input.to_string(); + + // Find all ${VAR_NAME} patterns + while let Some(start) = result.find("${") { + let end = result[start..].find('}').ok_or_else(|| { + anyhow::anyhow!("Unclosed environment variable substitution: {input}") + })?; + let end = start + end; + + let var_name = &result[start + 2..end]; + let var_value = std::env::var(var_name) + .with_context(|| format!("Environment variable not set: {var_name}"))?; + + result.replace_range(start..=end, &var_value); + } + + Ok(result) + } + + /// Check if Perplexity API key is configured. + pub fn has_perplexity_api_key(&self) -> bool { + self.perplexity_api_key + .as_ref() + .map(|k| !k.is_empty()) + .unwrap_or(false) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::env; + + #[test] + fn test_default_config() { + let config = ToolsConfig::default(); + assert!(config.perplexity_api_key.is_none()); + assert!(!config.has_perplexity_api_key()); + } + + #[test] + fn test_env_var_substitution() { + env::set_var("TEST_PERPLEXITY_KEY", "pplx-test-key"); + + let input = "${TEST_PERPLEXITY_KEY}"; + let result = ToolsConfig::substitute_env_var_in_string(input).unwrap(); + assert_eq!(result, "pplx-test-key"); + + env::remove_var("TEST_PERPLEXITY_KEY"); + } + + #[test] + fn test_has_perplexity_api_key() { + let mut config = ToolsConfig::default(); + assert!(!config.has_perplexity_api_key()); + + config.perplexity_api_key = Some(String::new()); + assert!(!config.has_perplexity_api_key()); + + config.perplexity_api_key = Some("pplx-xxx".to_string()); + assert!(config.has_perplexity_api_key()); + } +} diff --git a/crates/code_assistant/src/tools/core/dyn_tool.rs b/crates/code_assistant/src/tools/core/dyn_tool.rs index 8c86baf9..fc9f53a2 100644 --- a/crates/code_assistant/src/tools/core/dyn_tool.rs +++ b/crates/code_assistant/src/tools/core/dyn_tool.rs @@ -1,3 +1,4 @@ +use super::config::ToolsConfig; use super::render::Render; use super::result::ToolResult; use super::spec::ToolSpec; @@ -42,6 +43,9 @@ pub trait DynTool: Send + Sync + 'static { /// Get the static metadata for this tool fn spec(&self) -> ToolSpec; + /// Check if this tool is available based on configuration. + fn is_available(&self, config: &ToolsConfig) -> bool; + /// Invoke the tool with JSON parameters and get a type-erased output async fn invoke<'a>( &self, @@ -65,6 +69,10 @@ where Tool::spec(self) } + fn is_available(&self, config: &ToolsConfig) -> bool { + Tool::is_available(self, config) + } + async fn invoke<'a>( &self, context: &mut ToolContext<'a>, diff --git a/crates/code_assistant/src/tools/core/mod.rs b/crates/code_assistant/src/tools/core/mod.rs index 085de914..f0791000 100644 --- a/crates/code_assistant/src/tools/core/mod.rs +++ b/crates/code_assistant/src/tools/core/mod.rs @@ -1,4 +1,5 @@ // Core tools implementation +pub mod config; pub mod dyn_tool; pub mod registry; pub mod render; @@ -8,6 +9,7 @@ pub mod title; pub mod tool; // Re-export all core components for easier imports +pub use config::ToolsConfig; pub use dyn_tool::AnyOutput; pub use registry::ToolRegistry; pub use render::{Render, ResourcesTracker}; diff --git a/crates/code_assistant/src/tools/core/registry.rs b/crates/code_assistant/src/tools/core/registry.rs index ec5d47b4..f4a16117 100644 --- a/crates/code_assistant/src/tools/core/registry.rs +++ b/crates/code_assistant/src/tools/core/registry.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::sync::OnceLock; +use crate::tools::core::config::ToolsConfig; use crate::tools::core::dyn_tool::DynTool; use crate::tools::core::spec::ToolScope; use crate::tools::AnnotatedToolDefinition; @@ -29,9 +30,18 @@ impl ToolRegistry { } } - /// Register a tool in the registry + /// Register a tool in the registry. + /// The tool will only be registered if it's available based on the current configuration. pub fn register(&mut self, tool: Box) { - self.tools.insert(tool.spec().name.to_string(), tool); + let config = ToolsConfig::global(); + if tool.is_available(config) { + self.tools.insert(tool.spec().name.to_string(), tool); + } else { + tracing::debug!( + "Tool '{}' is not available (missing configuration)", + tool.spec().name + ); + } } /// Get a tool by name diff --git a/crates/code_assistant/src/tools/core/tool.rs b/crates/code_assistant/src/tools/core/tool.rs index 4b00994d..3dcbdd03 100644 --- a/crates/code_assistant/src/tools/core/tool.rs +++ b/crates/code_assistant/src/tools/core/tool.rs @@ -1,3 +1,4 @@ +use super::config::ToolsConfig; use super::render::Render; use super::result::ToolResult; use super::spec::ToolSpec; @@ -38,6 +39,14 @@ pub trait Tool: Send + Sync + 'static { /// Get the metadata for this tool fn spec(&self) -> ToolSpec; + /// Check if this tool is available based on configuration. + /// Tools that require external API keys or services should override this + /// to return false when their requirements are not met. + /// Default implementation returns true (tool is always available). + fn is_available(&self, _config: &ToolsConfig) -> bool { + true + } + /// Execute the tool with the given context and input /// The input may be modified during execution (e.g., for format-on-save) async fn execute<'a>( diff --git a/crates/code_assistant/src/tools/impls/perplexity_ask.rs b/crates/code_assistant/src/tools/impls/perplexity_ask.rs index 68f6dd33..464db22d 100644 --- a/crates/code_assistant/src/tools/impls/perplexity_ask.rs +++ b/crates/code_assistant/src/tools/impls/perplexity_ask.rs @@ -1,5 +1,5 @@ use crate::tools::core::{ - Render, ResourcesTracker, Tool, ToolContext, ToolResult, ToolScope, ToolSpec, + Render, ResourcesTracker, Tool, ToolContext, ToolResult, ToolScope, ToolSpec, ToolsConfig, }; use anyhow::Result; use serde::{Deserialize, Serialize}; @@ -123,13 +123,17 @@ impl Tool for PerplexityAskTool { } } + fn is_available(&self, config: &ToolsConfig) -> bool { + config.has_perplexity_api_key() + } + async fn execute<'a>( &self, _context: &mut ToolContext<'a>, input: &mut Self::Input, ) -> Result { - // Check if the API key exists - let api_key = std::env::var("PERPLEXITY_API_KEY").ok(); + // Get the API key from configuration + let api_key = ToolsConfig::global().perplexity_api_key.clone(); // Create a new Perplexity client let client = PerplexityClient::new(api_key); @@ -138,8 +142,7 @@ impl Tool for PerplexityAskTool { let query = input .messages .iter() - .filter(|m| m.role == "user") - .next_back() + .rfind(|m| m.role == "user") .map(|m| m.content.clone()) .unwrap_or_else(|| "No user query found".to_string()); diff --git a/crates/code_assistant/src/tools/impls/write_file.rs b/crates/code_assistant/src/tools/impls/write_file.rs index 9782934d..c5826b32 100644 --- a/crates/code_assistant/src/tools/impls/write_file.rs +++ b/crates/code_assistant/src/tools/impls/write_file.rs @@ -166,7 +166,7 @@ impl Tool for WriteFileTool { .write_file(&full_path, &input.content, input.append) .await { - Ok(mut full_content) => { + Ok(_) => { // If format-on-save applies, run the formatter if let Some(command_line) = project_config.format_command_for(&path) { let mut format_request = SandboxCommandRequest::default(); @@ -178,9 +178,8 @@ impl Tool for WriteFileTool { // Regardless of formatter success, try to re-read the file content if let Ok(updated) = explorer.read_file(&full_path).await { - full_content = updated; // Update the input content to the formatted content so the LLM sees it - input.content = full_content.clone(); + input.content = updated; } } diff --git a/crates/code_assistant/src/ui/gpui/elements.rs b/crates/code_assistant/src/ui/gpui/elements.rs index 7cea2bd1..ee5057ae 100644 --- a/crates/code_assistant/src/ui/gpui/elements.rs +++ b/crates/code_assistant/src/ui/gpui/elements.rs @@ -75,8 +75,20 @@ pub struct MessageContainer { /// each new request (see `UiEvent::StreamingStarted` in gpui/mod). When the user cancels /// streaming, all blocks that were created for that last, canceled request are removed. current_request_id: Arc>, + /// Current project for parameter filtering current_project: Arc>, + /// Tracks the last block type for hidden tool paragraph breaks + last_block_type_for_hidden_tool: Arc>>, + /// Flag indicating a hidden tool completed and we may need a paragraph break + needs_paragraph_break_after_hidden_tool: Arc>, +} + +/// Tracks the last block type for paragraph breaks after hidden tools +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum HiddenToolBlockType { + Text, + Thinking, } impl MessageContainer { @@ -86,6 +98,8 @@ impl MessageContainer { role, current_request_id: Arc::new(Mutex::new(0)), current_project: Arc::new(Mutex::new(String::new())), + last_block_type_for_hidden_tool: Arc::new(Mutex::new(None)), + needs_paragraph_break_after_hidden_tool: Arc::new(Mutex::new(false)), } } @@ -99,6 +113,34 @@ impl MessageContainer { *self.current_project.lock().unwrap() = project; } + /// Mark that a hidden tool completed - paragraph break may be needed before next text + pub fn mark_hidden_tool_completed(&self, _cx: &mut Context) { + *self.needs_paragraph_break_after_hidden_tool.lock().unwrap() = true; + } + + /// Check if we need a paragraph break after a hidden tool and return it if so + fn get_paragraph_break_if_needed( + &self, + current_block_type: HiddenToolBlockType, + ) -> Option { + let mut needs_break = self.needs_paragraph_break_after_hidden_tool.lock().unwrap(); + if !*needs_break { + return None; + } + + // Reset the flag + *needs_break = false; + + // Check if the block type matches the last one + let last_type = *self.last_block_type_for_hidden_tool.lock().unwrap(); + if last_type == Some(current_block_type) { + // Same type as before the hidden tool - need paragraph break + Some("\n\n".to_string()) + } else { + None + } + } + // Remove all blocks with the given request ID // Used when the user cancels a request while it is streaming pub fn remove_blocks_with_request_id(&self, request_id: u64, cx: &mut Context) { @@ -320,6 +362,12 @@ impl MessageContainer { pub fn add_or_append_to_text_block(&self, content: impl Into, cx: &mut Context) { self.finish_any_thinking_blocks(cx); + // Check if we need to insert a paragraph break after a hidden tool + let paragraph_prefix = self.get_paragraph_break_if_needed(HiddenToolBlockType::Text); + + // Track block type for future hidden tool events + *self.last_block_type_for_hidden_tool.lock().unwrap() = Some(HiddenToolBlockType::Text); + let content = content.into(); let mut elements = self.elements.lock().unwrap(); @@ -328,6 +376,9 @@ impl MessageContainer { last.update(cx, |view, cx| { if let Some(text_block) = view.block.as_text_mut() { + if let Some(prefix) = ¶graph_prefix { + text_block.content.push_str(prefix); + } text_block.content.push_str(&content); was_appended = true; cx.notify(); @@ -341,8 +392,13 @@ impl MessageContainer { // If we reach here, we need to add a new text block let request_id = *self.current_request_id.lock().unwrap(); + let final_content = if let Some(prefix) = paragraph_prefix { + format!("{}{}", prefix, content) + } else { + content.to_string() + }; let block = BlockData::TextBlock(TextBlock { - content: content.to_string(), + content: final_content, }); let view = cx.new(|cx| BlockView::new(block, request_id, self.current_project.clone(), cx)); elements.push(view); @@ -355,6 +411,12 @@ impl MessageContainer { content: impl Into, cx: &mut Context, ) { + // Check if we need to insert a paragraph break after a hidden tool + let paragraph_prefix = self.get_paragraph_break_if_needed(HiddenToolBlockType::Thinking); + + // Track block type for future hidden tool events + *self.last_block_type_for_hidden_tool.lock().unwrap() = Some(HiddenToolBlockType::Thinking); + let content = content.into(); let mut elements = self.elements.lock().unwrap(); @@ -363,6 +425,9 @@ impl MessageContainer { last.update(cx, |view, cx| { if let Some(thinking_block) = view.block.as_thinking_mut() { + if let Some(prefix) = ¶graph_prefix { + thinking_block.content.push_str(prefix); + } thinking_block.content.push_str(&content); was_appended = true; cx.notify(); @@ -376,7 +441,12 @@ impl MessageContainer { // If we reach here, we need to add a new thinking block let request_id = *self.current_request_id.lock().unwrap(); - let block = BlockData::ThinkingBlock(ThinkingBlock::new(content.to_string())); + let final_content = if let Some(prefix) = paragraph_prefix { + format!("{}{}", prefix, content) + } else { + content.to_string() + }; + let block = BlockData::ThinkingBlock(ThinkingBlock::new(final_content)); let view = cx.new(|cx| BlockView::new(block, request_id, self.current_project.clone(), cx)); elements.push(view); cx.notify(); diff --git a/crates/code_assistant/src/ui/gpui/mod.rs b/crates/code_assistant/src/ui/gpui/mod.rs index 85add068..605f15a6 100644 --- a/crates/code_assistant/src/ui/gpui/mod.rs +++ b/crates/code_assistant/src/ui/gpui/mod.rs @@ -570,11 +570,18 @@ impl Gpui { ); }); } + UiEvent::EndTool { id } => { self.update_all_messages(cx, |message_container, cx| { message_container.end_tool_use(&id, cx); }); } + UiEvent::HiddenToolCompleted => { + // Mark that a hidden tool completed - message container handles paragraph breaks + self.update_last_message(cx, |message, cx| { + message.mark_hidden_tool_completed(cx); + }); + } UiEvent::UpdatePlan { plan } => { if let Ok(mut plan_guard) = self.plan_state.lock() { @@ -1101,11 +1108,17 @@ impl Gpui { "GPUI: Tool {tool_id} attached terminal {terminal_id}; GUI terminal embedding unsupported" ); } + DisplayFragment::ReasoningComplete => { self.update_container(container, cx, |container, cx| { container.complete_reasoning(cx); }); } + DisplayFragment::HiddenToolCompleted => { + self.update_container(container, cx, |container, cx| { + container.mark_hidden_tool_completed(cx); + }); + } } } } @@ -1511,11 +1524,15 @@ impl UserInterface for Gpui { "GPUI: Tool {tool_id} attached terminal {terminal_id}; no dedicated UI hook" ); } + DisplayFragment::CompactionDivider { summary } => { self.push_event(UiEvent::DisplayCompactionSummary { summary: summary.clone(), }); } + DisplayFragment::HiddenToolCompleted => { + self.push_event(UiEvent::HiddenToolCompleted); + } } Ok(()) diff --git a/crates/code_assistant/src/ui/streaming/caret_processor.rs b/crates/code_assistant/src/ui/streaming/caret_processor.rs index 507d2980..791c055a 100644 --- a/crates/code_assistant/src/ui/streaming/caret_processor.rs +++ b/crates/code_assistant/src/ui/streaming/caret_processor.rs @@ -83,6 +83,7 @@ impl StreamProcessorTrait for CaretStreamProcessor { current_tool_id: String::new(), current_tool_name: String::new(), current_tool_hidden: false, + filter: Box::new(SmartToolFilter::new()), streaming_state: StreamingState::PreFirstTool, } @@ -759,18 +760,30 @@ impl CaretStreamProcessor { match &fragment { DisplayFragment::ToolName { .. } | DisplayFragment::ToolParameter { .. } - | DisplayFragment::ToolEnd { .. } | DisplayFragment::ToolOutput { .. } | DisplayFragment::ToolTerminal { .. } => { // Skip tool-related fragments for hidden tools return Ok(()); } + DisplayFragment::ToolEnd { .. } => { + // Emit HiddenToolCompleted so UI can handle paragraph breaks + return self + .emit_fragment_to_streaming_state(DisplayFragment::HiddenToolCompleted); + } _ => { // Allow non-tool fragments even when current tool is hidden } } } + self.emit_fragment_to_streaming_state(fragment) + } + + /// Send fragment to streaming state machine (handles buffering logic) + fn emit_fragment_to_streaming_state( + &mut self, + fragment: DisplayFragment, + ) -> Result<(), UIError> { match &self.streaming_state { StreamingState::Blocked => { // Already blocked, check if this is just whitespace and ignore silently @@ -914,6 +927,7 @@ impl CaretStreamProcessor { // Tool output - emit immediately (we've already decided to allow the tool) self.ui.display_fragment(&fragment)?; } + DisplayFragment::ReasoningComplete => { // Reasoning complete - buffer it until we know if next tool is allowed if let StreamingState::BufferingAfterTool { @@ -923,6 +937,10 @@ impl CaretStreamProcessor { buffered_fragments.push(fragment); } } + DisplayFragment::HiddenToolCompleted => { + // Hidden tool completed - emit immediately for UI to handle paragraph breaks + self.ui.display_fragment(&fragment)?; + } } } } diff --git a/crates/code_assistant/src/ui/streaming/caret_processor_tests.rs b/crates/code_assistant/src/ui/streaming/caret_processor_tests.rs index b696e6ae..8bfddb0b 100644 --- a/crates/code_assistant/src/ui/streaming/caret_processor_tests.rs +++ b/crates/code_assistant/src/ui/streaming/caret_processor_tests.rs @@ -884,3 +884,54 @@ fn test_smart_filter_blocks_write_tool_after_read_tool() { )); assert!(matches!(fragments[2], DisplayFragment::ToolEnd { .. })); } + +#[test] +fn test_hidden_tool_emits_hidden_tool_completed() { + use super::test_utils::print_fragments; + + let test_ui = TestUI::new(); + let ui_arc = Arc::new(test_ui.clone()); + let mut processor = CaretStreamProcessor::new(ui_arc, 42); + + // First send some text + processor + .process(&StreamingChunk::Text("Some text before.\n".to_string())) + .unwrap(); + + // Then process the hidden tool (update_plan is hidden) + processor + .process(&StreamingChunk::Text( + "^^^update_plan\nentries: [{\"content\": \"test\"}]\n^^^".to_string(), + )) + .unwrap(); + + // Then send more text (same type as before) + processor + .process(&StreamingChunk::Text("Next content.".to_string())) + .unwrap(); + + let fragments = test_ui.get_fragments(); + print_fragments(&fragments); + + // The tool fragments should be suppressed (update_plan is hidden) + assert!( + !fragments + .iter() + .any(|f| matches!(f, DisplayFragment::ToolName { .. })), + "ToolName should be suppressed for hidden tools" + ); + assert!( + !fragments + .iter() + .any(|f| matches!(f, DisplayFragment::ToolEnd { .. })), + "ToolEnd should be suppressed for hidden tools" + ); + + // Check that HiddenToolCompleted is emitted (UI handles paragraph breaks) + assert!( + fragments + .iter() + .any(|f| matches!(f, DisplayFragment::HiddenToolCompleted)), + "HiddenToolCompleted should be emitted for hidden tools" + ); +} diff --git a/crates/code_assistant/src/ui/streaming/json_processor.rs b/crates/code_assistant/src/ui/streaming/json_processor.rs index de5b8dd8..383a86d2 100644 --- a/crates/code_assistant/src/ui/streaming/json_processor.rs +++ b/crates/code_assistant/src/ui/streaming/json_processor.rs @@ -41,6 +41,7 @@ struct JsonProcessorState { in_thinking: bool, /// Track if we're at the beginning of a block (thinking/content) for text chunks at_block_start: bool, + /// True if any part of the current string value being parsed has been emitted. /// Used to detect and emit empty strings: "" emitted_string_part_for_current_key: bool, @@ -62,6 +63,7 @@ impl Default for JsonProcessorState { tool_name: String::new(), current_tool_hidden: false, in_thinking: false, + at_block_start: false, emitted_string_part_for_current_key: false, json_parsing_state: JsonParsingState::ExpectOpenBrace, @@ -82,16 +84,20 @@ pub struct JsonStreamProcessor { impl JsonStreamProcessor { /// Emit a fragment through this central function that handles hidden tool filtering - fn emit_fragment(&self, fragment: DisplayFragment) -> Result<(), UIError> { + fn emit_fragment(&mut self, fragment: DisplayFragment) -> Result<(), UIError> { // Filter out tool-related fragments for hidden tools if self.state.current_tool_hidden { match &fragment { - DisplayFragment::ToolName { .. } - | DisplayFragment::ToolParameter { .. } - | DisplayFragment::ToolEnd { .. } => { + DisplayFragment::ToolName { .. } | DisplayFragment::ToolParameter { .. } => { // Skip tool-related fragments for hidden tools return Ok(()); } + DisplayFragment::ToolEnd { .. } => { + // Emit HiddenToolCompleted so UI can handle paragraph breaks + return self + .ui + .display_fragment(&DisplayFragment::HiddenToolCompleted); + } _ => { // Allow non-tool fragments even when current tool is hidden } diff --git a/crates/code_assistant/src/ui/streaming/json_processor_tests.rs b/crates/code_assistant/src/ui/streaming/json_processor_tests.rs index fd2ec149..698ab704 100644 --- a/crates/code_assistant/src/ui/streaming/json_processor_tests.rs +++ b/crates/code_assistant/src/ui/streaming/json_processor_tests.rs @@ -1170,4 +1170,62 @@ mod tests { print_fragments(&fragments); assert_fragments_match(&expected_fragments, &fragments); } + + #[test] + fn test_hidden_tool_emits_hidden_tool_completed() { + let test_ui = TestUI::new(); + let ui_arc = Arc::new(test_ui.clone()); + let mut processor = JsonStreamProcessor::new(ui_arc, 42); + + // First send some text + processor + .process(&StreamingChunk::Text("Some text before.".to_string())) + .unwrap(); + + // Then process the hidden tool (update_plan is hidden) + processor + .process(&StreamingChunk::InputJson { + content: String::new(), + tool_name: Some("update_plan".to_string()), + tool_id: Some("tool-42-1".to_string()), + }) + .unwrap(); + processor + .process(&StreamingChunk::InputJson { + content: r#"{"entries":[{"content":"test"}]}"#.to_string(), + tool_name: None, + tool_id: None, + }) + .unwrap(); + + // Then send more text (same type as before) + processor + .process(&StreamingChunk::Text("Next content.".to_string())) + .unwrap(); + + let fragments = test_ui.get_fragments(); + print_fragments(&fragments); + + // The tool fragments should be suppressed (update_plan is hidden) + assert!( + !fragments + .iter() + .any(|f| matches!(f, DisplayFragment::ToolName { .. })), + "ToolName should be suppressed for hidden tools" + ); + assert!( + !fragments + .iter() + .any(|f| matches!(f, DisplayFragment::ToolEnd { .. })), + "ToolEnd should be suppressed for hidden tools" + ); + + // Check that HiddenToolCompleted is emitted (UI handles paragraph breaks) + assert!( + fragments + .iter() + .any(|f| matches!(f, DisplayFragment::HiddenToolCompleted)), + "HiddenToolCompleted should be emitted for hidden tools" + ); + } } diff --git a/crates/code_assistant/src/ui/streaming/mod.rs b/crates/code_assistant/src/ui/streaming/mod.rs index 906179e4..d2d4016f 100644 --- a/crates/code_assistant/src/ui/streaming/mod.rs +++ b/crates/code_assistant/src/ui/streaming/mod.rs @@ -53,6 +53,8 @@ pub enum DisplayFragment { ReasoningComplete, /// Divider indicating the conversation was compacted, with expandable summary text CompactionDivider { summary: String }, + /// A hidden tool completed - UI may need to insert paragraph break if next fragment is same type + HiddenToolCompleted, } /// Common trait for stream processors diff --git a/crates/code_assistant/src/ui/streaming/test_utils.rs b/crates/code_assistant/src/ui/streaming/test_utils.rs index 2213e969..5aa03087 100644 --- a/crates/code_assistant/src/ui/streaming/test_utils.rs +++ b/crates/code_assistant/src/ui/streaming/test_utils.rs @@ -168,9 +168,13 @@ pub fn print_fragments(fragments: &[DisplayFragment]) { terminal_id, } => println!(" [{i}] ToolTerminal(tool_id: {tool_id}, terminal_id: {terminal_id})"), DisplayFragment::ReasoningComplete => println!(" [{i}] ReasoningComplete"), + DisplayFragment::CompactionDivider { summary } => { println!(" [{i}] CompactionDivider: {summary}"); } + DisplayFragment::HiddenToolCompleted => { + println!(" [{i}] HiddenToolCompleted"); + } } } } diff --git a/crates/code_assistant/src/ui/streaming/xml_processor.rs b/crates/code_assistant/src/ui/streaming/xml_processor.rs index 38df8d52..187cbc24 100644 --- a/crates/code_assistant/src/ui/streaming/xml_processor.rs +++ b/crates/code_assistant/src/ui/streaming/xml_processor.rs @@ -74,6 +74,7 @@ impl StreamProcessorTrait for XmlStreamProcessor { state: ProcessorState::default(), ui, request_id, + filter: Box::new(SmartToolFilter::new()), streaming_state: StreamingState::PreFirstTool, } @@ -595,18 +596,29 @@ impl XmlStreamProcessor { // Filter out tool-related fragments for hidden tools if self.state.current_tool_hidden { match &fragment { - DisplayFragment::ToolName { .. } - | DisplayFragment::ToolParameter { .. } - | DisplayFragment::ToolEnd { .. } => { + DisplayFragment::ToolName { .. } | DisplayFragment::ToolParameter { .. } => { // Skip tool-related fragments for hidden tools return Ok(()); } + DisplayFragment::ToolEnd { .. } => { + // Emit HiddenToolCompleted so UI can handle paragraph breaks + return self + .emit_fragment_to_streaming_state(DisplayFragment::HiddenToolCompleted); + } _ => { // Allow non-tool fragments even when current tool is hidden } } } + self.emit_fragment_to_streaming_state(fragment) + } + + /// Send fragment to streaming state machine (handles buffering logic) + fn emit_fragment_to_streaming_state( + &mut self, + fragment: DisplayFragment, + ) -> Result<(), UIError> { match &self.streaming_state { StreamingState::Blocked => { // Already blocked, check if this is just whitespace and ignore silently @@ -750,6 +762,7 @@ impl XmlStreamProcessor { // Tool output - emit immediately (we've already decided to allow the tool) self.ui.display_fragment(&fragment)?; } + DisplayFragment::ReasoningComplete => { // Reasoning complete - buffer it if let StreamingState::BufferingAfterTool { @@ -759,6 +772,10 @@ impl XmlStreamProcessor { buffered_fragments.push(fragment); } } + DisplayFragment::HiddenToolCompleted => { + // Hidden tool completed - emit immediately for UI to handle paragraph breaks + self.ui.display_fragment(&fragment)?; + } } } } diff --git a/crates/code_assistant/src/ui/streaming/xml_processor_tests.rs b/crates/code_assistant/src/ui/streaming/xml_processor_tests.rs index 2bc1c0f3..606fd59c 100644 --- a/crates/code_assistant/src/ui/streaming/xml_processor_tests.rs +++ b/crates/code_assistant/src/ui/streaming/xml_processor_tests.rs @@ -977,10 +977,72 @@ mod tests { fragments[2], DisplayFragment::ToolParameter { .. } )); + assert!(matches!( fragments[3], DisplayFragment::ToolParameter { .. } )); assert!(matches!(fragments[4], DisplayFragment::ToolEnd { .. })); } + + #[test] + fn test_hidden_tool_emits_hidden_tool_completed() { + // Test that hidden tools (like update_plan) emit a paragraph break when suppressed + // AND the next fragment is the same type as the previous one + let input = concat!( + "Some text before.\n", + "\n", + "[{\"content\": \"test\"}]\n", + "", + "Text after hidden tool" + ); + + let test_ui = TestUI::new(); + let ui_arc = Arc::new(test_ui.clone()); + let mut processor = XmlStreamProcessor::new(ui_arc, 42); + + // Process the input + processor + .process(&StreamingChunk::Text(input.to_string())) + .unwrap(); + + let fragments = test_ui.get_fragments(); + print_fragments(&fragments); + + // The tool fragments should be suppressed (update_plan is hidden) + assert!( + !fragments + .iter() + .any(|f| matches!(f, DisplayFragment::ToolName { .. })), + "ToolName should be suppressed for hidden tools" + ); + assert!( + !fragments + .iter() + .any(|f| matches!(f, DisplayFragment::ToolEnd { .. })), + "ToolEnd should be suppressed for hidden tools" + ); + + // Check that HiddenToolCompleted is emitted (UI handles paragraph breaks) + assert!( + fragments + .iter() + .any(|f| matches!(f, DisplayFragment::HiddenToolCompleted)), + "HiddenToolCompleted should be emitted for hidden tools" + ); + + // Check that the text after the hidden tool is present + let combined_text: String = fragments + .iter() + .filter_map(|f| match f { + DisplayFragment::PlainText(text) => Some(text.as_str()), + _ => None, + }) + .collect(); + + assert!( + combined_text.contains("Text after hidden tool"), + "Should contain the text after the hidden tool" + ); + } } diff --git a/crates/code_assistant/src/ui/terminal/renderer.rs b/crates/code_assistant/src/ui/terminal/renderer.rs index 06eb9d9e..f6e4840f 100644 --- a/crates/code_assistant/src/ui/terminal/renderer.rs +++ b/crates/code_assistant/src/ui/terminal/renderer.rs @@ -92,10 +92,22 @@ pub struct TerminalRenderer { plan_expanded: bool, /// Last computed overflow (how many rows have been promoted so far); used to promote only deltas pub last_overflow: u16, + /// Maximum rows for input area (including 1 for content min + border) pub max_input_rows: u16, /// Spinner state for loading indication spinner_state: SpinnerState, + /// Tracks the last block type for hidden tool paragraph breaks + last_block_type_for_hidden_tool: Option, + /// Flag indicating a hidden tool completed and we may need a paragraph break + needs_paragraph_break_after_hidden_tool: bool, +} + +/// Tracks the last block type for paragraph breaks after hidden tools +#[derive(Debug, Clone, Copy, PartialEq)] +enum LastBlockType { + PlainText, + Thinking, } /// Type alias for the production terminal renderer @@ -115,11 +127,14 @@ impl TerminalRenderer { pending_user_message: None, current_error: None, info_message: None, + plan_state: None, plan_expanded: false, last_overflow: 0, max_input_rows: 5, // max input height (content lines + border line) spinner_state: SpinnerState::Hidden, + last_block_type_for_hidden_tool: None, + needs_paragraph_break_after_hidden_tool: false, }) } @@ -177,6 +192,43 @@ impl TerminalRenderer { /// If not, append a new block of that type /// Returns true if a new block was created, false if the last block was already the right type pub fn ensure_last_block_type(&mut self, block: MessageBlock) -> bool { + // Determine the block type for hidden tool tracking + let current_block_type = match &block { + MessageBlock::PlainText(_) => Some(LastBlockType::PlainText), + MessageBlock::Thinking(_) => Some(LastBlockType::Thinking), + _ => None, + }; + + // Check if we need to insert a paragraph break after a hidden tool + if self.needs_paragraph_break_after_hidden_tool { + if let (Some(last_type), Some(current_type)) = + (self.last_block_type_for_hidden_tool, current_block_type) + { + if last_type == current_type { + // Same type as before the hidden tool - insert paragraph break + if let Some(live_message) = self.live_message.as_mut() { + if let Some(last_block) = live_message.blocks.last_mut() { + match last_block { + MessageBlock::PlainText(text_block) => { + text_block.content.push_str("\n\n"); + } + MessageBlock::Thinking(thinking_block) => { + thinking_block.content.push_str("\n\n"); + } + _ => {} + } + } + } + } + } + self.needs_paragraph_break_after_hidden_tool = false; + } + + // Track the block type for future hidden tool events + if let Some(block_type) = current_block_type { + self.last_block_type_for_hidden_tool = Some(block_type); + } + let live_message = self.live_message.as_mut() .expect("ensure_last_block_type called without an active live message - call start_new_message first"); @@ -195,6 +247,11 @@ impl TerminalRenderer { needs_new_block } + /// Mark that a hidden tool completed - paragraph break may be needed before next text + pub fn mark_hidden_tool_completed(&mut self) { + self.needs_paragraph_break_after_hidden_tool = true; + } + /// Set or unset a pending user message (displayed while streaming) pub fn set_pending_user_message(&mut self, message: Option) { self.pending_user_message = message; diff --git a/crates/code_assistant/src/ui/terminal/ui.rs b/crates/code_assistant/src/ui/terminal/ui.rs index 79fa5555..be3ef547 100644 --- a/crates/code_assistant/src/ui/terminal/ui.rs +++ b/crates/code_assistant/src/ui/terminal/ui.rs @@ -271,11 +271,19 @@ impl UserInterface for TerminalTuiUI { renderer_guard.add_or_update_tool_parameter(&tool_id, name, value); } } + UiEvent::EndTool { id: _ } => { // EndTool just marks the end of parameter streaming // The actual status comes later via UpdateToolStatus // For now, we don't change the status here - wait for UpdateToolStatus } + UiEvent::HiddenToolCompleted => { + // Mark that a hidden tool completed - renderer handles paragraph breaks + if let Some(renderer) = self.renderer.lock().await.as_ref() { + let mut renderer_guard = renderer.lock().await; + renderer_guard.mark_hidden_tool_completed(); + } + } UiEvent::StreamingStopped { id, cancelled, @@ -461,9 +469,14 @@ impl UserInterface for TerminalTuiUI { summary: summary.clone(), }); } + DisplayFragment::ReasoningComplete => { // For terminal UI, no specific action needed for reasoning completion } + DisplayFragment::HiddenToolCompleted => { + // Signal that a hidden tool completed - renderer handles paragraph breaks + self.push_event(UiEvent::HiddenToolCompleted); + } } Ok(()) diff --git a/crates/code_assistant/src/ui/ui_events.rs b/crates/code_assistant/src/ui/ui_events.rs index 6c716973..6415e1bd 100644 --- a/crates/code_assistant/src/ui/ui_events.rs +++ b/crates/code_assistant/src/ui/ui_events.rs @@ -51,8 +51,11 @@ pub enum UiEvent { message: Option, output: Option, }, + /// End a tool invocation EndTool { id: String }, + /// A hidden tool completed - UI may need paragraph break before next text + HiddenToolCompleted, /// Add an image to the message AddImage { media_type: String, data: String }, /// Append streaming tool output diff --git a/tools.example.json b/tools.example.json new file mode 100644 index 00000000..db054131 --- /dev/null +++ b/tools.example.json @@ -0,0 +1,3 @@ +{ + "perplexity_api_key": "${PERPLEXITY_API_KEY}" +}