diff --git a/crates/forge_domain/src/chat_request.rs b/crates/forge_domain/src/chat_request.rs new file mode 100644 index 0000000000000000000000000000000000000000..2b9b5824945fac8c6f2ff7ae7aeabd44fe0a6ac4 --- /dev/null +++ b/crates/forge_domain/src/chat_request.rs @@ -0,0 +1,17 @@ +use derive_setters::Setters; +use serde::{Deserialize, Serialize}; + +use crate::{ConversationId, Event}; + +#[derive(Debug, Serialize, Deserialize, Clone, Setters)] +#[setters(into, strip_option)] +pub struct ChatRequest { + pub event: Event, + pub conversation_id: ConversationId, +} + +impl ChatRequest { + pub fn new(content: Event, conversation_id: ConversationId) -> Self { + Self { event: content, conversation_id } + } +} diff --git a/crates/forge_domain/src/command.rs b/crates/forge_domain/src/command.rs new file mode 100644 index 0000000000000000000000000000000000000000..4f909309167854f61ce9938cfc9adb28c513076a --- /dev/null +++ b/crates/forge_domain/src/command.rs @@ -0,0 +1,21 @@ +use derive_setters::Setters; +use serde::Deserialize; + +/// A user-defined command loaded from a Markdown file with YAML frontmatter. +/// +/// Commands are discovered from `.md` files in the forge commands directories +/// and made available as slash commands in the UI. The `name` and `description` +/// come from YAML frontmatter; the `prompt` is the Markdown body of the file. +#[derive(Debug, Clone, Default, Deserialize, Setters, PartialEq)] +#[setters(into, strip_option)] +pub struct Command { + /// The command name used to invoke it (e.g. `github-pr-description`). + #[serde(default)] + pub name: String, + /// Short description shown in the command list. + #[serde(default)] + pub description: String, + /// The prompt template body (Markdown content after the frontmatter). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt: Option, +} diff --git a/crates/forge_domain/src/console.rs b/crates/forge_domain/src/console.rs new file mode 100644 index 0000000000000000000000000000000000000000..dc057310b40f273a47533d38ba9385411bf34d81 --- /dev/null +++ b/crates/forge_domain/src/console.rs @@ -0,0 +1,15 @@ +use std::io; + +/// Trait for synchronized output writing. +/// Provides two output channels (primary and error) with flush support. +/// Implementors must ensure thread-safe writes. +pub trait ConsoleWriter: Send + Sync { + /// Writes bytes to primary output. + fn write(&self, buf: &[u8]) -> io::Result; + /// Writes bytes to error output. + fn write_err(&self, buf: &[u8]) -> io::Result; + /// Flushes primary output. + fn flush(&self) -> io::Result<()>; + /// Flushes error output. + fn flush_err(&self) -> io::Result<()>; +} diff --git a/crates/forge_domain/src/conversation.rs b/crates/forge_domain/src/conversation.rs new file mode 100644 index 0000000000000000000000000000000000000000..c0bde6e4e85a67176fbc242b3198eb5711cca1eb --- /dev/null +++ b/crates/forge_domain/src/conversation.rs @@ -0,0 +1,358 @@ +use std::str::FromStr; + +use chrono::{DateTime, Utc}; +use derive_more::derive::Display; +use derive_setters::Setters; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::{Context, Error, Metrics, Result, TokenCount}; + +#[derive(Debug, Default, Display, Serialize, Deserialize, Clone, PartialEq, Eq, Hash)] +#[serde(transparent)] +pub struct ConversationId(Uuid); + +impl Copy for ConversationId {} + +impl ConversationId { + pub fn generate() -> Self { + Self(Uuid::new_v4()) + } + + pub fn into_string(&self) -> String { + self.0.to_string() + } + + pub fn parse(value: impl ToString) -> Result { + Ok(Self( + Uuid::parse_str(&value.to_string()).map_err(Error::ConversationId)?, + )) + } +} + +impl FromStr for ConversationId { + type Err = Error; + + fn from_str(s: &str) -> Result { + Self::parse(s) + } +} + +#[derive(Debug, Setters, Serialize, Deserialize, Clone)] +#[setters(into)] +pub struct Conversation { + pub id: ConversationId, + pub title: Option, + pub context: Option, + pub metrics: Metrics, + pub metadata: MetaData, +} + +#[derive(Debug, Setters, Serialize, Deserialize, Clone)] +#[setters(into)] +pub struct MetaData { + pub created_at: DateTime, + pub updated_at: Option>, +} + +impl MetaData { + pub fn new(created_at: DateTime) -> Self { + Self { created_at, updated_at: None } + } +} + +impl Conversation { + pub fn new(id: ConversationId) -> Self { + let created_at = Utc::now(); + let metrics = Metrics::default().started_at(created_at); + Self { + id, + metrics, + metadata: MetaData::new(created_at), + title: None, + context: None, + } + } + /// Creates a new conversation with a new conversation ID. + /// + /// This is a convenience constructor that automatically generates a unique + /// conversation ID, making it easy to create new conversations without + /// having to manually create the ID. + pub fn generate() -> Self { + Self::new(ConversationId::generate()) + } + + /// Generates an HTML representation of the conversation + /// + /// This method uses Handlebars to render the conversation as HTML + /// from the template file, including all agents, events, and variables. + /// + /// # Errors + /// - If the template file cannot be found or read + /// - If the Handlebars template registration fails + /// - If the template rendering fails + pub fn to_html(&self) -> String { + // Instead of using Handlebars, we now use our Element DSL + crate::conversation_html::render_conversation_html(self) + } + + /// Generates an HTML representation with related agent conversations + /// + /// Creates a single HTML document containing the main conversation + /// and all related agent conversations with anchor links for navigation. + pub fn to_html_with_related(&self, related: &[Conversation]) -> String { + crate::conversation_html::render_conversation_html_with_related(self, related) + } + + /// Returns a vector of user messages, selecting the first message from + /// each consecutive sequence of user messages. + pub fn first_user_messages(&self) -> Vec<&crate::ContextMessage> { + self.context + .as_ref() + .map(|ctx| ctx.first_user_messages()) + .unwrap_or_default() + } + + /// Returns the total token usage across all messages in the conversation. + /// + /// This is a convenience method that aggregates usage from the context, + /// if available. + pub fn accumulated_usage(&self) -> Option { + self.context.as_ref().and_then(|ctx| ctx.accumulate_usage()) + } + + pub fn usage(&self) -> Option { + self.context + .iter() + .flat_map(|ctx| ctx.messages.iter()) + .flat_map(|msg| msg.usage.into_iter()) + .last() + } + + pub fn token_count(&self) -> Option { + self.context.as_ref().map(|ctx| ctx.token_count()) + } + + pub fn accumulated_cost(&self) -> Option { + self.accumulated_usage().and_then(|usage| usage.cost) + } + + /// Calculates the total cost including related conversations. + /// + /// This method sums the costs of the current conversation and all related + /// conversations provided. It's useful for computing the total cost of a + /// conversation tree when agent tools spawn child conversations. + /// + /// # Arguments + /// + /// * `conversations` - A slice of conversations to include in cost + /// calculation + pub fn total_cost(conversations: &[Conversation]) -> Option { + let costs: Vec = conversations + .iter() + .filter_map(|conv| conv.accumulated_cost()) + .collect(); + + if costs.is_empty() { + return None; + } + + Some(costs.iter().sum()) + } + + /// Returns the number of messages in the conversation context. + /// + /// Returns `0` if the context has not been initialized yet. + pub fn len(&self) -> usize { + self.context + .as_ref() + .map(|ctx| ctx.messages.len()) + .unwrap_or(0) + } + + /// Returns `true` if the conversation context has no messages. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Extracts all related conversation IDs from agent tool calls. + /// + /// This method scans through all tool results in the conversation's context + /// and collects conversation IDs from AI tool values, which are created + /// when agent tools are called and trigger new conversations. + pub fn related_conversation_ids(&self) -> Vec { + self.context + .as_ref() + .map(|ctx| { + ctx.messages + .iter() + .filter_map(|msg| msg.as_tool_result()) + .flat_map(|result| &result.output.values) + .filter_map(|value| { + if let crate::ToolValue::AI { conversation_id, .. } = value { + Some(conversation_id) + } else { + None + } + }) + .copied() + .collect() + }) + .unwrap_or_default() + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + use crate::{Context, ContextMessage, ToolOutput, ToolResult, ToolValue}; + + #[test] + fn test_related_conversation_ids_empty() { + let conversation = Conversation::generate(); + let actual = conversation.related_conversation_ids(); + assert_eq!(actual, vec![]); + } + + #[test] + fn test_related_conversation_ids_with_agent_calls() { + let agent_conv_id_1 = ConversationId::generate(); + let agent_conv_id_2 = ConversationId::generate(); + + let context = Context::default() + .add_message(ContextMessage::user("Test task", None)) + .add_message(ContextMessage::assistant("Working on it", None, None, None)) + .add_message(ContextMessage::Tool(ToolResult::new("agent_tool").output( + Ok(ToolOutput { + is_error: false, + values: vec![ToolValue::AI { + value: "Agent result".to_string(), + conversation_id: agent_conv_id_1, + }], + }), + ))) + .add_message(ContextMessage::assistant("Continuing", None, None, None)) + .add_message(ContextMessage::Tool( + ToolResult::new("another_agent").output(Ok(ToolOutput { + is_error: false, + values: vec![ToolValue::AI { + value: "Another agent result".to_string(), + conversation_id: agent_conv_id_2, + }], + })), + )); + + let conversation = Conversation::generate().context(context); + let actual = conversation.related_conversation_ids(); + + assert_eq!(actual.len(), 2); + assert!(actual.contains(&agent_conv_id_1)); + assert!(actual.contains(&agent_conv_id_2)); + } + + #[test] + fn test_related_conversation_ids_mixed_tool_results() { + let agent_conv_id = ConversationId::generate(); + + let context = Context::default() + .add_message(ContextMessage::user("Test task", None)) + .add_message(ContextMessage::Tool( + ToolResult::new("regular_tool").output(Ok(ToolOutput::text("Regular result"))), + )) + .add_message(ContextMessage::Tool(ToolResult::new("agent_tool").output( + Ok(ToolOutput { + is_error: false, + values: vec![ToolValue::AI { + value: "Agent result".to_string(), + conversation_id: agent_conv_id, + }], + }), + ))) + .add_message(ContextMessage::Tool( + ToolResult::new("another_regular_tool") + .output(Ok(ToolOutput::text("Another regular result"))), + )); + + let conversation = Conversation::generate().context(context); + let actual = conversation.related_conversation_ids(); + + assert_eq!(actual, vec![agent_conv_id]); + } + + #[test] + fn test_total_cost() { + use crate::{MessageEntry, Usage}; + + // Create main conversation with cost + // Create main conversation with cost + let main_usage = Usage { cost: Some(0.01), ..Usage::default() }; + let main_entry: MessageEntry = + ContextMessage::assistant("Response", None, None, None).into(); + let main_context = Context::default() + .add_message(ContextMessage::user("Test", None)) + .add_entry(main_entry.usage(main_usage)); + + let main_conv = Conversation::generate().context(main_context); + + // Create related conversations with costs + let related_usage_1 = Usage { cost: Some(0.02), ..Usage::default() }; + let related_entry_1: MessageEntry = + ContextMessage::assistant("Result 1", None, None, None).into(); + let related_context_1 = Context::default() + .add_message(ContextMessage::user("Task 1", None)) + .add_entry(related_entry_1.usage(related_usage_1)); + + let related_conv_1 = Conversation::generate().context(related_context_1); + + let related_usage_2 = Usage { cost: Some(0.03), ..Usage::default() }; + let related_entry_2: MessageEntry = + ContextMessage::assistant("Result 2", None, None, None).into(); + let related_context_2 = Context::default() + .add_message(ContextMessage::user("Task 2", None)) + .add_entry(related_entry_2.usage(related_usage_2)); + let related_conv_2 = Conversation::generate().context(related_context_2); + + let actual = Conversation::total_cost(&[main_conv, related_conv_1, related_conv_2]); + + // Check that cost is approximately 0.06 (accounting for floating point + // precision) + assert!(actual.is_some()); + assert!((actual.unwrap() - 0.06).abs() < 0.0001); + } + + #[test] + fn test_total_cost_no_costs() { + let main_conv = Conversation::generate(); + let related_conv = Conversation::generate(); + + let actual = Conversation::total_cost(&[main_conv, related_conv]); + let expected = None; + + assert_eq!(actual, expected); + } + + #[test] + fn test_total_cost_partial_costs() { + use crate::{MessageEntry, Usage}; + + // Main conversation has no cost + let main_conv = Conversation::generate(); + + // Related conversation has cost + let related_usage = Usage { cost: Some(0.05), ..Usage::default() }; + let related_entry: MessageEntry = + ContextMessage::assistant("Result", None, None, None).into(); + let related_context = Context::default() + .add_message(ContextMessage::user("Task", None)) + .add_entry(related_entry.usage(related_usage)); + + let related_conv = Conversation::generate().context(related_context); + + let actual = Conversation::total_cost(&[main_conv, related_conv]); + let expected = Some(0.05); + + assert_eq!(actual, expected); + } +} diff --git a/crates/forge_domain/src/conversation_html.rs b/crates/forge_domain/src/conversation_html.rs new file mode 100644 index 0000000000000000000000000000000000000000..aec2a1e9051c4233600243685764e388c3641a7d --- /dev/null +++ b/crates/forge_domain/src/conversation_html.rs @@ -0,0 +1,513 @@ +use forge_template::Element; +use serde_json::to_string_pretty; + +use crate::context::ContextMessage; +use crate::conversation::Conversation; + +/// Renders a conversation as an HTML document +/// +/// Creates a complete HTML page displaying the conversation's information +/// including: +/// - Basic information (ID, title) +/// - Reasoning configuration +/// - Usage statistics (token counts and costs) +/// - Context messages with tool calls and reasoning details +/// - Available tools +/// - Error styling for tool call failures +/// +/// # Arguments +/// +/// * `conversation` - The conversation to render +pub fn render_conversation_html(conversation: &Conversation) -> String { + let c_title = format!( + "Title: {}", + conversation + .title + .clone() + .unwrap_or(conversation.id.to_string()) + ); + let html = Element::new("html") + .attr("lang", "en") + .append( + Element::new("head") + .append(Element::new("meta").attr("charset", "UTF-8")) + .append( + Element::new("meta") + .attr("name", "viewport") + .attr("content", "width=device-width, initial-scale=1.0"), + ) + .append(Element::new("title").text(&c_title)) + .append(Element::new("style").text(include_str!("conversation_style.css"))), // Includes tool-call-error styles + ) + .append( + Element::new("body") + // Combined Information Table + .append(create_info_table(conversation)) + // Conversation Context Section + .append(create_conversation_context_section(conversation)) + // Tools Section + .append(create_tools_section(conversation)), + ); + + format!("\n{}", html.render()) +} + +/// Renders a conversation with related agent conversations in a single HTML +/// document +/// +/// Creates a complete HTML page with the main conversation followed by +/// related agent conversations. Uses anchor links for navigation. +/// +/// # Arguments +/// +/// * `conversation` - The main conversation to render +/// * `related` - Related agent conversations to include +pub fn render_conversation_html_with_related( + conversation: &Conversation, + related: &[Conversation], +) -> String { + let c_title = format!( + "Title: {}", + conversation + .title + .clone() + .unwrap_or(conversation.id.to_string()) + ); + + let mut body = Element::new("body") + // Combined Information Table + .append(create_info_table(conversation)) + // Conversation Context Section + .append(create_conversation_context_section(conversation)) + // Tools Section + .append(create_tools_section(conversation)); + + // Add related conversations section + if !related.is_empty() { + body = body.append(Element::new("div.section").append( + Element::new("h2").text(format!("Related Agent Conversations ({})", related.len())), + )); + + for related_conv in related { + let anchor_id = format!("conversation-{}", related_conv.id); + body = body.append( + Element::new("div.related-conversation") + .attr("id", &anchor_id) + .append( + Element::new("div.back-to-main").append( + Element::new("a") + .attr("href", "#") + .text("⬆ Back to main conversation"), + ), + ) + .append(create_info_table(related_conv)) + .append(create_conversation_context_section(related_conv)) + .append(create_tools_section(related_conv)), + ); + } + } + + let html = Element::new("html") + .attr("lang", "en") + .append( + Element::new("head") + .append(Element::new("meta").attr("charset", "UTF-8")) + .append( + Element::new("meta") + .attr("name", "viewport") + .attr("content", "width=device-width, initial-scale=1.0"), + ) + .append(Element::new("title").text(&c_title)) + .append(Element::new("style").text(include_str!("conversation_style.css"))), + ) + .append(body); + + format!("\n{}", html.render()) +} + +/// Creates a table row with a label and value +fn create_table_row(label: impl Into, value: impl Into) -> Element { + Element::new("tr") + .append(Element::new("th").text(label.into())) + .append(Element::new("td").text(value.into())) +} + +/// Creates a combined information table with all conversation metadata +fn create_info_table(conversation: &Conversation) -> Element { + let section = Element::new("div.section").append(Element::new("h2").text("Conversation")); + + let mut table = Element::new("table") + .append(create_table_row("ID", conversation.id.to_string())) + .append(create_table_row( + "Title", + conversation + .title + .clone() + .unwrap_or_else(|| "No title".to_string()), + )); + + // Add reasoning configuration if available + if let Some(context) = &conversation.context { + if let Some(reasoning_config) = &context.reasoning { + let status = match reasoning_config.enabled { + Some(true) => "Enabled", + Some(false) => "Disabled", + None => "Not specified", + }; + table = table + .append(create_table_row("Reasoning Status", status)) + .append(create_table_row( + "Reasoning Effort", + format!("{:?}", reasoning_config.effort), + )); + + if let Some(max_tokens) = reasoning_config.max_tokens { + table = table.append(create_table_row( + "Reasoning Max Tokens", + format!("{max_tokens:?}"), + )); + } + } + + if let Some(max_tokens) = context.max_tokens { + table = table.append(create_table_row( + "Max Output Tokens", + format!("{:?}", max_tokens), + )) + } + + // Add usage information if available + if let Some(usage) = context.accumulate_usage() { + let cache_percentage = if *usage.prompt_tokens > 0 { + (*usage.cached_tokens as f64 / *usage.prompt_tokens as f64 * 100.0) as usize + } else { + 0 + }; + + let cached_display = if cache_percentage > 0 { + format!("{} [{}%]", usage.cached_tokens, cache_percentage) + } else { + format!("{}", usage.cached_tokens) + }; + + table = table + .append(create_table_row( + "Input Tokens", + format!("{}", usage.prompt_tokens), + )) + .append(create_table_row("Cached Tokens", cached_display)) + .append(create_table_row( + "Output Tokens", + format!("{}", usage.completion_tokens), + )) + .append(create_table_row( + "Total Tokens", + format!("{}", usage.total_tokens), + )); + + if let Some(cost) = usage.cost { + table = table.append(create_table_row("Cost", format!("${:.4}", cost))); + } + } + } + + section.append(table) +} + +/// Creates a tools section displaying all available tools +fn create_tools_section(conversation: &Conversation) -> Element { + let section = Element::new("div.section").append(Element::new("h2").text("Tools")); + + if let Some(context) = &conversation.context { + if !context.tools.is_empty() { + let tools_elm = + Element::new("div.tools-section").append(context.tools.iter().map(|tool| { + Element::new("details.message-card.message-tool") + .append( + Element::new("summary").append(Element::span(tool.name.to_string())), + ) + .append( + Element::new("div.main-content") + .append(Element::new("p").append( + Element::new("strong").text(tool.description.to_string()), + )) + .append(Element::new("pre").text( + to_string_pretty(&tool.input_schema).unwrap_or_default(), + )), + ) + })); + section.append(tools_elm) + } else { + section.append(Element::new("p").text("No tools available")) + } + } else { + section.append(Element::new("p").text("No tools available")) + } +} + +/// Creates a usage information section for a message +fn create_message_usage_section(usage: &crate::message::Usage) -> Element { + let cache_percentage = if *usage.prompt_tokens > 0 { + (*usage.cached_tokens as f64 / *usage.prompt_tokens as f64 * 100.0) as usize + } else { + 0 + }; + + let cached_display = if cache_percentage > 0 { + format!("{} [{}%]", usage.cached_tokens, cache_percentage) + } else { + format!("{}", usage.cached_tokens) + }; + + let mut usage_div = Element::new("span") + .append(Element::new("strong").text("📊 Usage {")) + .append( + Element::new("span") + .append( + Element::new("span.usage-item").text(format!("input: {}", usage.prompt_tokens)), + ) + .append(Element::new("span.usage-item").text(format!("cached: {}", cached_display))) + .append( + Element::new("span.usage-item") + .text(format!("output: {}", usage.completion_tokens)), + ) + .append( + Element::new("span.usage-item").text(format!("total: {}", usage.total_tokens)), + ), + ) + .append(Element::new("strong").text("}")); + + if let Some(cost) = usage.cost { + usage_div = usage_div + .append(Element::new("span.usage-item.usage-cost").text(format!("Cost: ${:.4}", cost))); + } + + usage_div +} + +fn create_conversation_context_section(conversation: &Conversation) -> Element { + let section = Element::new("div.section").append(Element::new("h2").text("Messages")); + + // Add context if available + if let Some(context) = &conversation.context { + let context_messages = Element::new("div.context-section").append( + context.messages.iter().map(|message_entry| { + match &**message_entry { + ContextMessage::Text(content_message) => { + // Convert role to lowercase for the class + let role_lowercase = content_message.role.to_string().to_lowercase(); + + let mut header = + Element::new("summary").text(format!("{}", content_message.role)); + + if let Some(model) = &content_message.model { + header = header + .append(Element::new("strong").text(" 🤖 model:")) + .append(Element::new("span").text(model)); + } + + // Add usage information + if let Some(usage) = &message_entry.usage { + header = header.append(create_message_usage_section(usage)) + } + + // Add reasoning indicator if reasoning details are present + let has_reasoning = content_message.reasoning_details.as_ref().is_some_and(|d| !d.is_empty()) + || content_message.thought_signature.is_some(); + if has_reasoning + { + header = header.append( + Element::new("span.reasoning-indicator").text(" 🧠 Reasoning"), + ); + } + + let message_elm = + Element::new(format!("details.message-card.message-{role_lowercase}")) + .append(header); + + // Add thought signature + let mut message_elm = if let Some(sig) = &content_message.thought_signature { + message_elm.append( + Element::new("div.thought-signature") + .append(Element::new("strong").text("💭 Thought Signature: ")) + .append(Element::new("pre").text(sig)), + ) + } else { + message_elm + }; + + // Add reasoning details + message_elm = if let Some(reasoning_details) = + &content_message.reasoning_details + { + if !reasoning_details.is_empty() { + message_elm.append(Element::new("div.reasoning-section").append( + reasoning_details.iter().map(|reasoning_detail| { + if let Some(text) = &reasoning_detail.text { + Element::new("div.reasoning-content") + .append( + Element::new("strong").text("🧠 Reasoning: "), + ) + .append(Element::new("pre").text(text)) + } else { + Element::new("div") + } + }), + )) + } else { + message_elm + } + } else { + message_elm + }; + + // Add main content + let message_elm = message_elm.append( + Element::new("div.main-content") + .append(Element::new("pre").text(&content_message.content)), + ); + + // Add tool calls if any + + if let Some(tool_calls) = &content_message.tool_calls { + if !tool_calls.is_empty() { + message_elm.append(Element::new("div").append( + tool_calls.iter().map(|tool_call| { + Element::new("div.tool-call") + .append( + Element::new("p").append( + Element::new("strong") + .text(tool_call.name.to_string()), + ), + ) + .append(tool_call.call_id.as_ref().map(|call_id| { + Element::new("p") + .append(Element::new("strong").text("ID: ")) + .text(call_id.as_str()) + })) + .append( + Element::new("p").append( + Element::new("strong").text("Arguments: "), + ), + ) + .append( + Element::new("pre").text( + to_string_pretty(&tool_call.arguments) + .unwrap_or_default(), + ), + ) + }), + )) + } else { + message_elm + } + } else { + message_elm + } + } + ContextMessage::Tool(tool_result) => { + // Tool Message - apply error styling if the tool result is an error + let message_class = if tool_result.output.is_error { + "details.message-card.message-tool.tool-call-error" + } else { + "details.message-card.message-tool" + }; + + Element::new(message_class) + .append( + Element::new("summary") + .append(Element::new("strong").text("Tool Result: ")) + .append(Element::span(tool_result.name.as_str())), + ) + .append(Element::new("div.main-content").append( + tool_result.output.values.iter().filter_map(|value| { + match value { + crate::ToolValue::Text(text) => Some( + Element::new("div") + .append(Element::new("pre").text(text)), + ), + crate::ToolValue::Image(image) => { + Some(Element::new("img").attr("src", image.url())) + } + crate::ToolValue::Empty => None, + crate::ToolValue::AI { value, conversation_id } => { + // Use anchor link to navigate within the same HTML + let anchor_id = format!("conversation-{}", conversation_id); + Some( + Element::new("div.agent-conversation") + .append( + Element::new("p") + .append(Element::new("strong").text("🤖 Agent Conversation: ")) + .append( + Element::new("a") + .attr("href", format!("#{}", anchor_id)) + .attr("title", "Click to view agent conversation") + .text(conversation_id.to_string()) + ) + ) + .append(Element::new("pre").text(value)), + ) + }, + } + }), + )) + } + ContextMessage::Image(image) => { + // Image message + Element::new("div.message-card.message-user") + .append(Element::new("strong").text("Image Attachment")) + .append(Element::new("img").attr("src", image.url())) + } + } + }), + ); + + // Create tool choice section if available + let context_elm = if let Some(tool_choice) = &context.tool_choice { + context_messages + .append(Element::new("strong").text("Tool Choice")) + .append(Element::new("div.tool-choice").append( + Element::new("pre").text(to_string_pretty(tool_choice).unwrap_or_default()), + )) + } else { + context_messages + }; + + // Add temperature if available + let context_elm = if let Some(temperature) = context.temperature { + context_elm.append( + Element::new("p") + .append(Element::new("strong").text("Temperature: ")) + .text(format!("{temperature}")), + ) + } else { + context_elm + }; + + section.append(context_elm) + } else { + section.append(Element::new("p").text("No context available")) + } +} + +#[cfg(test)] +mod tests { + use forge_test_kit::json_fixture; + + use super::*; + + #[tokio::test] + async fn test_render_conversation_html_snapshot() { + // Load the conversation from the fixture file + let conversation: Conversation = json_fixture!("tests/fixtures/conversation.json").await; + + // Render the HTML + let html = render_conversation_html(&conversation); + + // Convert HTML string to bytes for binary snapshot + let html_bytes = html.into_bytes(); + + // Binary snapshot with exact .html extension + insta::assert_binary_snapshot!("conversation.html", html_bytes); + } +} diff --git a/crates/forge_domain/src/conversation_style.css b/crates/forge_domain/src/conversation_style.css new file mode 100644 index 0000000000000000000000000000000000000000..5aea1b1aa027c2671c8aaf120e4fd4a0544c9882 --- /dev/null +++ b/crates/forge_domain/src/conversation_style.css @@ -0,0 +1,313 @@ +body { + font-family: "JetBrainsMono Nerd Font", "JetBrains Mono", "Fira Code", "Cascadia Code", "Consolas", "Menlo", "Monaco", monospace; + line-height: 1.6; + color: #333; + font-size: 1em; + max-width: 1200px; + margin: 0 auto; +} + +summary { + cursor: pointer; +} + +h1, +h2, +h3 { + color: #2c3e50; +} + +h1 { + border-bottom: 2px solid #eee; + padding-bottom: 10px; +} + +.section { + padding: 20px; + background-color: #f9f9f9; + border: 1px solid #eee; +} + +.agent, +.event, +.variable { + padding: 15px; + background-color: white; + border: 1px solid #ddd; +} + +.agent-header, +.event-header { + display: flex; + justify-content: space-between; + padding-bottom: 5px; + border-bottom: 1px solid #eee; +} + +pre { + padding: 10px; + overflow-x: auto; +} + +code, pre { + font-family: "JetBrainsMono Nerd Font", "JetBrains Mono", "Fira Code", "Cascadia Code", "Consolas", "Menlo", "Monaco", monospace; +} + +table { + width: 100%; + border-collapse: collapse; +} + +th, +td { + padding: 8px; + text-align: left; +} + +th { + background-color: #f2f2f2; + font-weight: 600; + width: 30%; +} + +td { + background-color: #ffffff; +} + +.section table { + margin-top: 10px; + border: 1px solid #ddd; +} + +.section table th { + border-bottom: 1px solid #ddd; +} + +.section table tr:not(:last-child) td { + border-bottom: 1px solid #eee; +} + +.message-card { + padding: 15px; + border: 1px solid #ddd; +} + +.message-system { + background-color: #f8f9fa; +} + +.message-user { + background-color: #e9f5ff; +} + +.message-assistant { + background-color: #f0f7e6; +} + +.message-tool { + background-color: #fff8e6; +} + +.tool-call, +.tool-result { + background-color: #f5f5f5; +} + +.tool-choice { + padding: 5px 10px; + background-color: #eef; + display: inline-block; +} + +img { + max-width: 100%; +} + +/* Reasoning styles */ +.reasoning-section { + margin: 12px 0; + padding: 16px; + background-color: #f0f4f8; + border-left: 4px solid #7c3aed; +} + +.reasoning-content { + margin-bottom: 12px; +} + +.reasoning-content strong { + color: #7c3aed; + font-weight: 600; +} + +.reasoning-content pre { + margin-top: 8px; + padding: 12px; + background-color: #faf5ff; + border: 1px solid #e9d5ff; + white-space: pre-wrap; + word-wrap: break-word; + font-size: 0.95em; + color: #6b7280; + font-style: italic; + line-height: 1.6; +} + +.main-content { + margin: 12px 0; +} + +.main-content strong { + color: #1f2937; + font-weight: 600; +} + +.main-content pre { + margin-top: 8px; + padding: 12px; + background-color: #ffffff; + border: 1px solid #e5e7eb; + white-space: pre-wrap; + word-wrap: break-word; + font-size: 0.95em; + color: #1f2937; + line-height: 1.6; +} + +.reasoning-indicator { + margin-left: 8px; + color: #7c3aed; + font-size: 0.9em; +} + +/* Usage Section Styles */ +.usage-section { + background-color: #f8f9fa; + border-left: 3px solid #6c757d; + padding: 8px 12px; + margin-top: 12px; + font-size: 0.9em; +} + +.usage-section strong { + display: block; + margin-bottom: 6px; + color: #495057; +} + +.usage-details { + display: flex; + flex-wrap: wrap; + gap: 12px; +} + +.usage-item { + color: #6c757d; + font-weight: 500; +} + +.usage-cost { + color: #28a745; + font-weight: 600; +} + +/* Tool Call Error Styles */ +.tool-call-error { + background-color: #fff5f5; + border-left: 4px solid #ef4444; + padding: 16px; + margin: 12px 0; + border-radius: 4px; +} + +.tool-call-error .error-cause { + color: #dc2626; + font-weight: 600; + margin-bottom: 8px; +} + +.tool-call-error .error-reflection { + background-color: #fee2e2; + padding: 12px; + border-radius: 4px; + margin-top: 8px; + font-size: 0.9em; + color: #7f1d1d; +} + +/* Agent Conversation Link Styles */ +.agent-conversation { + background-color: #f0f9ff; + border-left: 4px solid #3b82f6; + padding: 16px; + margin: 12px 0; + border-radius: 4px; +} + +.agent-conversation p { + margin: 0 0 12px 0; +} + +.agent-conversation strong { + color: #1e40af; + font-weight: 600; +} + +.agent-conversation a { + color: #2563eb; + text-decoration: none; + font-weight: 600; + padding: 2px 6px; + background-color: #dbeafe; + border-radius: 3px; + transition: background-color 0.2s; +} + +.agent-conversation a:hover { + background-color: #bfdbfe; + text-decoration: underline; +} + +.agent-conversation pre { + margin-top: 8px; + padding: 12px; + background-color: #ffffff; + border: 1px solid #e0f2fe; + white-space: pre-wrap; + word-wrap: break-word; + font-size: 0.95em; + color: #1e293b; + line-height: 1.6; +} + +/* Related Conversation Styles */ +.related-conversation { + margin-top: 40px; + padding-top: 20px; + border-top: 3px solid #3b82f6; + scroll-margin-top: 20px; +} + +.back-to-main { + margin-bottom: 20px; + padding: 10px; + background-color: #f0f9ff; + border-radius: 4px; + text-align: center; +} + +.back-to-main a { + color: #2563eb; + text-decoration: none; + font-weight: 600; + padding: 8px 16px; + background-color: #dbeafe; + border-radius: 4px; + display: inline-block; + transition: background-color 0.2s; +} + +.back-to-main a:hover { + background-color: #bfdbfe; + text-decoration: none; +} \ No newline at end of file diff --git a/crates/forge_domain/src/data_gen.rs b/crates/forge_domain/src/data_gen.rs new file mode 100644 index 0000000000000000000000000000000000000000..0ce3ba735d109406ad74922cdbc9c6f25ee06207 --- /dev/null +++ b/crates/forge_domain/src/data_gen.rs @@ -0,0 +1,28 @@ +use std::path::PathBuf; + +use derive_setters::Setters; +use serde::{Deserialize, Serialize}; + +/// Parameters for data generation operations +/// +/// This struct encapsulates the configuration parameters needed for generating +/// data in various contexts. It provides control over the amount of data to +/// generate, formatting options, and other generation-specific settings. +#[derive(Debug, Clone, Serialize, Deserialize, Setters, PartialEq, fake::Dummy)] +#[setters(into, strip_option)] +pub struct DataGenerationParameters { + /// Path to input JSONL file for data generation + pub input: PathBuf, + + /// Path to JSON schema file for LLM tool definition + pub schema: PathBuf, + + /// Path to Handlebars template file for system prompt + pub system_prompt: Option, + + /// Path to Handlebars template file for user prompt + pub user_prompt: Option, + + /// Maximum number of concurrent LLM requests + pub concurrency: usize, +} diff --git a/crates/forge_domain/src/env.rs b/crates/forge_domain/src/env.rs new file mode 100644 index 0000000000000000000000000000000000000000..6f2d0ab0763cc8d6aa09aba30ece401fc54b146b --- /dev/null +++ b/crates/forge_domain/src/env.rs @@ -0,0 +1,388 @@ +use std::hash::{DefaultHasher, Hash, Hasher}; +use std::path::PathBuf; + +use derive_more::Display; +use derive_setters::Setters; +use serde::{Deserialize, Serialize}; + +use crate::{Effort, ModelConfig}; + +/// All discrete mutations that can be applied to the application configuration. +/// +/// Instead of replacing the entire config, callers describe exactly which field +/// they want to change. Implementations receive a list of operations, apply +/// each in order, and persist the result atomically. +#[derive(Debug, Clone, PartialEq)] +pub enum ConfigOperation { + /// Set the active session provider and model atomically. + /// + /// When the provider differs from the current session provider the entire + /// session (provider + model) is replaced atomically. When they match only + /// the model field is updated. + SetSessionConfig(ModelConfig), + /// Set the commit-message generation configuration. + /// + /// `None` clears the commit configuration so the active session + /// provider/model is used for commit message generation. + SetCommitConfig(Option), + /// Set the shell-command suggestion configuration. + SetSuggestConfig(ModelConfig), + /// Set the reasoning effort level for all agents. + SetReasoningEffort(Effort), +} + +const VERSION: &str = match option_env!("APP_VERSION") { + Some(val) => val, + None => env!("CARGO_PKG_VERSION"), +}; + +/// Represents the minimal runtime environment in which the application is +/// running. +/// +/// Contains only the five fields that cannot be sourced from [`ForgeConfig`]: +/// `os`, `cwd`, `home`, `shell`, and `base_path`. All configuration +/// values previously carried here are now accessed through +/// `EnvironmentInfra::get_config()`. +#[derive(Debug, Setters, Clone, PartialEq, Serialize, Deserialize, fake::Dummy)] +#[serde(rename_all = "camelCase")] +#[setters(strip_option)] +pub struct Environment { + /// The operating system of the environment. + pub os: String, + /// The current working directory. + pub cwd: PathBuf, + /// The home directory. + pub home: Option, + /// The shell being used. + pub shell: String, + /// The base path relative to which everything else is stored. + pub base_path: PathBuf, +} + +impl Environment { + pub fn log_path(&self) -> PathBuf { + self.base_path.join("logs") + } + + /// Returns the history file path. + /// + /// # Arguments + /// * `custom_path` - An optional custom path sourced from + /// `ForgeConfig::custom_history_path`. When present it overrides the + /// default location inside `base_path`. + pub fn history_path(&self, custom_path: Option<&PathBuf>) -> PathBuf { + custom_path + .cloned() + .unwrap_or(self.base_path.join(".forge_history")) + } + pub fn snapshot_path(&self) -> PathBuf { + self.base_path.join("snapshots") + } + pub fn mcp_user_config(&self) -> PathBuf { + self.base_path.join(".mcp.json") + } + + /// Returns the path where the MCP trust store is persisted across restarts. + pub fn mcp_trust_path(&self) -> PathBuf { + self.base_path.join(".mcp_trust.json") + } + + pub fn agent_path(&self) -> PathBuf { + self.base_path.join("agents") + } + pub fn agent_cwd_path(&self) -> PathBuf { + self.cwd.join(".forge/agents") + } + + pub fn permissions_path(&self) -> PathBuf { + self.base_path.join("permissions.yaml") + } + + pub fn mcp_local_config(&self) -> PathBuf { + self.cwd.join(".mcp.json") + } + + pub fn version(&self) -> String { + VERSION.to_string() + } + + pub fn app_config(&self) -> PathBuf { + self.base_path.join(".config.json") + } + + pub fn database_path(&self) -> PathBuf { + self.base_path.join(".forge.db") + } + + /// Returns the path to the cache directory + pub fn cache_dir(&self) -> PathBuf { + self.base_path.join("cache") + } + + /// Returns the global skills directory path (~/forge/skills) + pub fn global_skills_path(&self) -> PathBuf { + self.base_path.join("skills") + } + + /// Returns the agents skills directory path (~/.agents/skills) + /// + /// Returns `None` when the home directory cannot be determined. + pub fn agents_skills_path(&self) -> Option { + self.home.as_ref().map(|home| home.join(".agents/skills")) + } + + /// Returns the project-local skills directory path (.forge/skills) + pub fn local_skills_path(&self) -> PathBuf { + self.cwd.join(".forge/skills") + } + + /// Returns the global commands directory path (base_path/commands) + pub fn command_path(&self) -> PathBuf { + self.base_path.join("commands") + } + + /// Returns the project-local commands directory path (.forge/commands) + pub fn command_path_local(&self) -> PathBuf { + self.cwd.join(".forge/commands") + } + + /// Returns the global AGENTS.md path (base_path/AGENTS.md) + pub fn global_agentsmd_path(&self) -> PathBuf { + self.base_path.join("AGENTS.md") + } + + /// Returns the project-local AGENTS.md path (cwd/AGENTS.md) + pub fn local_agentsmd_path(&self) -> PathBuf { + self.cwd.join("AGENTS.md") + } + + /// Returns the plans directory path relative to the current working + /// directory (cwd/plans) + pub fn plans_path(&self) -> PathBuf { + self.cwd.join("plans") + } + + /// Returns the path to the custom provider configuration file + /// (base_path/provider.json) + pub fn provider_config_path(&self) -> PathBuf { + self.base_path.join("provider.json") + } + + /// Returns the path to the credentials file where provider API keys are + /// stored + pub fn credentials_path(&self) -> PathBuf { + self.base_path.join(".credentials.json") + } + + pub fn workspace_hash(&self) -> WorkspaceHash { + let mut hasher = DefaultHasher::default(); + self.cwd.hash(&mut hasher); + + WorkspaceHash(hasher.finish()) + } +} + +#[derive(Clone, Copy, Display)] +pub struct WorkspaceHash(u64); +impl WorkspaceHash { + pub fn new(id: u64) -> Self { + WorkspaceHash(id) + } + + pub fn id(&self) -> u64 { + self.0 + } +} + +#[cfg(test)] +mod tests { + use fake::{Fake, Faker}; + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_agent_cwd_path() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.cwd(PathBuf::from("/current/working/dir")); + + let actual = fixture.agent_cwd_path(); + let expected = PathBuf::from("/current/working/dir/.forge/agents"); + + assert_eq!(actual, expected); + } + + #[test] + fn test_agent_cwd_path_independent_from_agent_path() { + let fixture: Environment = Faker.fake(); + let fixture = fixture + .cwd(PathBuf::from("/different/current/dir")) + .base_path(PathBuf::from("/completely/different/base")); + + let agent_path = fixture.agent_path(); + let agent_cwd_path = fixture.agent_cwd_path(); + let expected_agent_path = PathBuf::from("/completely/different/base/agents"); + let expected_agent_cwd_path = PathBuf::from("/different/current/dir/.forge/agents"); + + // Verify that agent_path uses base_path + assert_eq!(agent_path, expected_agent_path); + + // Verify that agent_cwd_path is independent and always relative to CWD + assert_eq!(agent_cwd_path, expected_agent_cwd_path); + + // Verify they are different paths + assert_ne!(agent_path, agent_cwd_path); + } + + #[test] + fn test_global_skills_path() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.base_path(PathBuf::from("/home/user/.forge")); + + let actual = fixture.global_skills_path(); + let expected = PathBuf::from("/home/user/.forge/skills"); + + assert_eq!(actual, expected); + } + + #[test] + fn test_agents_skills_path_with_home() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.home(PathBuf::from("/home/user")); + + let actual = fixture.agents_skills_path(); + let expected = Some(PathBuf::from("/home/user/.agents/skills")); + + assert_eq!(actual, expected); + } + + #[test] + fn test_agents_skills_path_without_home() { + let fixture: Environment = Faker.fake(); + // Explicitly clear the home field + let mut fixture = fixture; + fixture.home = None; + + let actual = fixture.agents_skills_path(); + + assert_eq!(actual, None); + } + + #[test] + fn test_local_skills_path() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.cwd(PathBuf::from("/projects/my-app")); + + let actual = fixture.local_skills_path(); + let expected = PathBuf::from("/projects/my-app/.forge/skills"); + + assert_eq!(actual, expected); + } + + #[test] + fn test_skills_paths_independent() { + let fixture: Environment = Faker.fake(); + let fixture = fixture + .cwd(PathBuf::from("/projects/my-app")) + .base_path(PathBuf::from("/home/user/.forge")); + + let global_path = fixture.global_skills_path(); + let local_path = fixture.local_skills_path(); + + let expected_global = PathBuf::from("/home/user/.forge/skills"); + let expected_local = PathBuf::from("/projects/my-app/.forge/skills"); + + // Verify global path uses base_path + assert_eq!(global_path, expected_global); + + // Verify local path uses cwd + assert_eq!(local_path, expected_local); + + // Verify they are different paths + assert_ne!(global_path, local_path); + } + + #[test] + fn test_command_path() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.base_path(PathBuf::from("/home/user/.forge")); + + let actual = fixture.command_path(); + let expected = PathBuf::from("/home/user/.forge/commands"); + + assert_eq!(actual, expected); + } + + #[test] + fn test_command_path_local() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.cwd(PathBuf::from("/projects/my-app")); + + let actual = fixture.command_path_local(); + let expected = PathBuf::from("/projects/my-app/.forge/commands"); + + assert_eq!(actual, expected); + } + + #[test] + fn test_command_paths_independent() { + let fixture: Environment = Faker.fake(); + let fixture = fixture + .cwd(PathBuf::from("/projects/my-app")) + .base_path(PathBuf::from("/home/user/.forge")); + + let global_path = fixture.command_path(); + let local_path = fixture.command_path_local(); + + let expected_global = PathBuf::from("/home/user/.forge/commands"); + let expected_local = PathBuf::from("/projects/my-app/.forge/commands"); + + assert_eq!(global_path, expected_global); + assert_eq!(local_path, expected_local); + assert_ne!(global_path, local_path); + } + + #[test] + fn test_global_agents_md_path() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.base_path(PathBuf::from("/home/user/.forge")); + + let actual = fixture.global_agentsmd_path(); + let expected = PathBuf::from("/home/user/.forge/AGENTS.md"); + + assert_eq!(actual, expected); + } + + #[test] + fn test_local_agents_md_path() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.cwd(PathBuf::from("/projects/my-app")); + + let actual = fixture.local_agentsmd_path(); + let expected = PathBuf::from("/projects/my-app/AGENTS.md"); + + assert_eq!(actual, expected); + } + + #[test] + fn test_plans_path() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.cwd(PathBuf::from("/projects/my-app")); + + let actual = fixture.plans_path(); + let expected = PathBuf::from("/projects/my-app/plans"); + + assert_eq!(actual, expected); + } + + #[test] + fn test_provider_config_path() { + let fixture: Environment = Faker.fake(); + let fixture = fixture.base_path(PathBuf::from("/home/user/.forge")); + + let actual = fixture.provider_config_path(); + let expected = PathBuf::from("/home/user/.forge/provider.json"); + + assert_eq!(actual, expected); + } +} diff --git a/crates/forge_domain/src/error.rs b/crates/forge_domain/src/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..31db4e4203da7c4d748126d2fff754c3714e6989 --- /dev/null +++ b/crates/forge_domain/src/error.rs @@ -0,0 +1,178 @@ +use std::pin::Pin; + +use derive_more::From; +use forge_json_repair::JsonRepairError; +use thiserror::Error; + +use crate::{AgentId, ConversationId, ProviderId, WorkspaceId}; + +// NOTE: Deriving From for error is a really bad idea. This is because you end +// up converting errors incorrectly without much context. For eg: You don't want +// all serde error to be treated as the same. Instead we want to know exactly +// where that serde failure happened and for what kind of value. +#[derive(Debug, Error, From)] +pub enum Error { + #[error("Missing tool name")] + ToolCallMissingName, + + #[error("Missing tool id")] + ToolCallMissingId, + + #[error("Unsupported role: {0}")] + #[from(skip)] + UnsupportedRole(String), + + #[error("{0}")] + EToolCallArgument(ToolCallArgumentError), + + #[error("JSON deserialization error: {error}")] + #[from(skip)] + ToolCallArgument { + error: JsonRepairError, + args: String, + }, + + #[error("JSON deserialization error: {error}")] + #[from(skip)] + AgentCallArgument { error: serde_json::error::Error }, + + #[error("Invalid tool call XML: {0}")] + #[from(skip)] + ToolCallParse(String), + + #[error("Invalid conversation id: {0}")] + ConversationId(uuid::Error), + + #[error("Agent not found in the arena: {0}")] + AgentUndefined(AgentId), + + #[error("Variable not found in output: {0}")] + #[from(skip)] + UndefinedVariable(String), + + #[error("Head agent not found")] + HeadAgentUndefined, + + #[error("Agent '{0}' has reached max turns of {1}")] + MaxTurnsReached(AgentId, u64), + + #[error("Conversation with ID '{0}' not found")] + ConversationNotFound(ConversationId), + + #[error("Missing description for agent: {0}")] + #[from(skip)] + MissingAgentDescription(AgentId), + + #[error("Missing model for agent: {0}")] + #[from(skip)] + MissingModel(AgentId), + + #[error("No model defined for agent: {0}")] + #[from(skip)] + NoModelDefined(AgentId), + + #[error("Empty completion received - no content, tool calls, or valid finish reason")] + EmptyCompletion, + + #[error( + "The model refused to generate a response (safety/content filter). \ + Retrying the same request will produce the same refusal - rephrase \ + the request or switch to a different model." + )] + Refusal, + + #[error(transparent)] + Retryable(anyhow::Error), + + #[error("Environment variable {env_var} not found for provider {provider}")] + EnvironmentVariableNotFound { + provider: ProviderId, + env_var: String, + }, + + #[error("Provider {provider} is not available. Login again to configure it.")] + ProviderNotAvailable { provider: ProviderId }, + + #[error("Failed to create VertexAI provider: {message}")] + VertexAiConfiguration { message: String }, + + // Indexing errors + #[error("No indexing authentication found")] + AuthTokenNotFound, + + #[error("Workspace not found")] + WorkspaceNotFound, + + #[error("Workspace already initialized with id: {0}")] + WorkspaceAlreadyInitialized(WorkspaceId), + + #[error("Failed to sync {count} file(s)")] + SyncFailed { count: usize }, + + #[error("No default provider and model configured.")] + NoDefaultSession, +} + +pub type Result = std::result::Result; +pub type BoxStream = + Pin> + Send>>; + +pub type ResultStream = std::result::Result, E>; + +#[derive(Debug, derive_more::From)] +pub struct ToolCallArgumentError(eserde::DeserializationErrors); + +impl std::error::Error for ToolCallArgumentError {} + +impl std::fmt::Display for ToolCallArgumentError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "Invalid tool call arguments:")?; + for error in self.0.iter() { + writeln!(f, "- {error}")?; + } + Ok(()) + } +} + +impl Error { + pub fn into_retryable(self) -> Self { + use anyhow::anyhow; + Self::Retryable(anyhow!(self)) + } + + pub fn env_var_not_found(provider: ProviderId, env_var: &str) -> Self { + Self::EnvironmentVariableNotFound { provider, env_var: env_var.to_string() } + } + + pub fn provider_not_available(provider: ProviderId) -> Self { + Self::ProviderNotAvailable { provider } + } + + pub fn vertex_ai_config(message: impl Into) -> Self { + Self::VertexAiConfiguration { message: message.into() } + } + + pub fn sync_failed(count: usize) -> Self { + Self::SyncFailed { count } + } +} + +#[cfg(test)] +mod test { + use forge_json_repair::JsonRepairError; + use serde_json::Value; + + use crate::Error; + + #[test] + fn test_debug_serde_error() { + let args = "{a: 1}"; + let serde_error = serde_json::from_str::(args).unwrap_err(); + let a = Error::ToolCallArgument { + error: JsonRepairError::from(serde_error), + args: args.to_string(), + }; + let a = anyhow::anyhow!(a); + eprintln!("{:?}", a.root_cause()); + } +} diff --git a/crates/forge_domain/src/event.rs b/crates/forge_domain/src/event.rs new file mode 100644 index 0000000000000000000000000000000000000000..4609aaf76d2d508b04214352c7bc9d3acdfa99c4 --- /dev/null +++ b/crates/forge_domain/src/event.rs @@ -0,0 +1,227 @@ +use std::collections::HashMap; + +use derive_more::{Deref, From}; +use derive_setters::Setters; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::{Attachment, NamedTool, Template, TerminalContext, ToolName}; + +/// Represents a partial event structure used for CLI event dispatching +/// +/// This is an intermediate structure for parsing event JSON from the CLI +/// before converting it to a full Event type. +#[derive(Debug, Default, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct UserCommand { + pub name: String, + pub template: Template, + pub parameters: Vec, +} + +impl UserCommand { + pub fn new>>( + name: impl ToString, + value: V, + parameters: Vec, + ) -> Self { + Self { name: name.to_string(), template: value.into(), parameters } + } +} + +impl From for Event { + fn from(value: UserCommand) -> Self { + Event::new(EventValue::Command(value)) + } +} + +impl> From for EventValue { + fn from(value: T) -> Self { + EventValue::Text(UserPrompt(value.as_ref().to_owned())) + } +} + +// We'll use simple strings for JSON schema compatibility +#[derive(Debug, Deserialize, Serialize, Clone, Setters)] +#[setters(into, strip_option)] +pub struct Event { + pub id: String, + pub value: Option, + pub timestamp: String, + pub attachments: Vec, + + /// Contains additional context about the prompt that should typically be + /// included after the `value` as a user message. + pub additional_context: Option, +} + +#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)] +pub enum EventValue { + Text(UserPrompt), + Command(UserCommand), +} + +impl EventValue { + pub fn as_user_prompt(&self) -> Option<&UserPrompt> { + match self { + EventValue::Text(user_prompt) => Some(user_prompt), + EventValue::Command(_) => None, + } + } + + pub fn as_command(&self) -> Option<&UserCommand> { + match self { + EventValue::Text(_user_prompt) => None, + EventValue::Command(user_command) => Some(user_command), + } + } + + pub fn text(str: impl ToString) -> Self { + EventValue::Text(UserPrompt(str.to_string())) + } +} + +#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq, From, Deref)] +#[serde(transparent)] +pub struct UserPrompt(String); + +#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Setters)] +pub struct EventContext { + event: EventContextValue, + suggestions: Vec, + variables: HashMap, + current_date: String, + /// Structured terminal context injected by [`TerminalContextService`], + /// or `None` when terminal context is unavailable or disabled. + #[serde(default, skip_serializing_if = "Option::is_none")] + terminal_context: Option, +} + +#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Setters)] +pub struct EventContextValue { + pub name: String, + pub value: String, +} + +impl EventContextValue { + pub fn new>(value: S) -> Self { + Self { name: String::new(), value: value.into() } + } +} + +impl EventContext { + pub fn new(event: impl Into) -> Self { + Self { + event: event.into(), + suggestions: Default::default(), + variables: Default::default(), + current_date: chrono::Local::now().format("%Y-%m-%d").to_string(), + terminal_context: None, + } + } + + /// Converts this EventContext into a feedback event by setting the event + /// name to "feedback". This should be used when the context already + /// contains user messages. + pub fn into_feedback(mut self) -> Self { + self.event.name = "feedback".to_string(); + self + } + + /// Converts this EventContext into a new task event by setting the event + /// name to "task". This should be used when this is a new task without + /// prior user messages. + pub fn into_task(mut self) -> Self { + self.event.name = "task".to_string(); + self + } +} + +impl NamedTool for Event { + fn tool_name() -> ToolName { + ToolName::new("forge_tool_event_dispatch") + } +} + +impl Event { + pub fn new>(value: V) -> Self { + let id = uuid::Uuid::new_v4().to_string(); + let timestamp = chrono::Utc::now().to_rfc3339(); + + Self { + id, + value: Some(value.into()), + timestamp, + attachments: Vec::new(), + additional_context: None, + } + } + + pub fn empty() -> Self { + let id = uuid::Uuid::new_v4().to_string(); + let timestamp = chrono::Utc::now().to_rfc3339(); + + Self { + id, + value: None, + timestamp, + attachments: Vec::new(), + additional_context: None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_into_feedback() { + let event = EventContextValue::new(""); + let context = EventContext::new(event); + + let feedback_context = context.into_feedback(); + + assert_eq!(feedback_context.event.name, "feedback"); + } + + #[test] + fn test_into_task() { + let event = EventContextValue::new(""); + let context = EventContext::new(event); + + let task_context = context.into_task(); + + assert_eq!(task_context.event.name, "task"); + } + + #[test] + fn test_into_feedback_idempotent() { + let event = EventContextValue::new(""); + let context = EventContext::new(event); + + // Call into_feedback twice + let feedback_context = context.into_feedback().into_feedback(); + + assert_eq!(feedback_context.event.name, "feedback"); + } + + #[test] + fn test_into_task_idempotent() { + let event = EventContextValue::new(""); + let context = EventContext::new(event); + + // Call into_task twice + let task_context = context.into_task().into_task(); + + assert_eq!(task_context.event.name, "task"); + } + + #[test] + fn test_chaining_methods() { + let event = EventContextValue::new("initial content"); + let context = EventContext::new(event).into_task(); + + assert_eq!(context.event.name, "task"); + assert_eq!(context.event.value, "initial content"); + } +} diff --git a/crates/forge_domain/src/file_operation.rs b/crates/forge_domain/src/file_operation.rs new file mode 100644 index 0000000000000000000000000000000000000000..ee32e441e1260d04c6f6eb252a0806cf8fde89e8 --- /dev/null +++ b/crates/forge_domain/src/file_operation.rs @@ -0,0 +1,25 @@ +use derive_setters::Setters; +use serde::{Deserialize, Serialize}; + +use crate::ToolKind; + +/// Tracks metrics for individual file changes +#[derive(Debug, Clone, PartialEq, Setters, Serialize, Deserialize)] +#[setters(into)] +pub struct FileOperation { + pub lines_added: u64, + pub lines_removed: u64, + /// Content hash of the file. None if file is unreadable (deleted, no + /// permissions, etc.) + pub content_hash: Option, + /// The tool that performed this operation + pub tool: ToolKind, +} + +impl FileOperation { + /// Creates a new FileChangeMetrics with the specified tool + /// Other fields default to zero/None and can be set using setters + pub fn new(tool: ToolKind) -> Self { + Self { lines_added: 0, lines_removed: 0, content_hash: None, tool } + } +} diff --git a/crates/forge_domain/src/group_by_key.rs b/crates/forge_domain/src/group_by_key.rs new file mode 100644 index 0000000000000000000000000000000000000000..50a1bdf634564d657c090c9e385d563e07a0bc5d --- /dev/null +++ b/crates/forge_domain/src/group_by_key.rs @@ -0,0 +1,28 @@ +use std::collections::HashMap; +use std::hash::Hash; + +/// Trait for grouping a collection of items by a key extracted from each item. +pub trait GroupByKey { + fn group_by_key(value: Vec, key_fn: impl Fn(&V) -> K) -> HashMap>; +} + +impl GroupByKey for Vec { + fn group_by_key(values: Vec, key_fn: impl Fn(&V) -> K) -> HashMap> { + values + .iter() + .fold( + HashMap::>::new(), + |mut acc: HashMap<_, _>, chunk| { + acc.entry(key_fn(chunk)).or_default().push(chunk); + acc + }, + ) + .into_iter() + .map(|(key, mut chunks)| { + chunks.sort(); + + (key.clone(), chunks.into_iter().cloned().collect::>()) + }) + .collect::>>() + } +} diff --git a/crates/forge_domain/src/hook.rs b/crates/forge_domain/src/hook.rs new file mode 100644 index 0000000000000000000000000000000000000000..47579d7a43c3e6d37e8c947600ae6d7fa3dbe4a5 --- /dev/null +++ b/crates/forge_domain/src/hook.rs @@ -0,0 +1,1136 @@ +use async_trait::async_trait; +use derive_more::From; +use derive_setters::Setters; + +use crate::{Agent, ChatCompletionMessageFull, Conversation, ModelId, ToolCallFull, ToolResult}; + +/// A container for lifecycle events with agent and model ID context +/// +/// This struct provides a consistent structure for all lifecycle events, +/// containing the agent and model ID along with event-specific payload data. +#[derive(Debug, PartialEq, Clone)] +pub struct EventData { + /// The agent associated with this event + pub agent: Agent, + /// The model ID being used + pub model_id: ModelId, + /// Event-specific payload data + pub payload: P, +} + +impl EventData

{ + /// Creates a new event with the given agent, model ID, and payload + pub fn new(agent: Agent, model_id: ModelId, payload: P) -> Self { + Self { agent, model_id, payload } + } +} + +/// Payload for the Start event +#[derive(Debug, PartialEq, Clone, Default)] +pub struct StartPayload; + +/// Payload for the End event +#[derive(Debug, PartialEq, Clone, Default)] +pub struct EndPayload; + +/// Payload for the Request event +#[derive(Debug, PartialEq, Clone, Setters)] +#[setters(into)] +pub struct RequestPayload { + /// The number of requests made + pub request_count: usize, +} + +impl RequestPayload { + /// Creates a new request payload + pub fn new(request_count: usize) -> Self { + Self { request_count } + } +} + +/// Payload for the Response event +#[derive(Debug, PartialEq, Clone, Setters)] +#[setters(into)] +pub struct ResponsePayload { + /// The full response message from the LLM + pub message: ChatCompletionMessageFull, +} + +impl ResponsePayload { + /// Creates a new response payload + pub fn new(message: ChatCompletionMessageFull) -> Self { + Self { message } + } +} + +/// Payload for the ToolcallStart event +#[derive(Debug, PartialEq, Clone, Setters)] +#[setters(into)] +pub struct ToolcallStartPayload { + /// The tool call details + pub tool_call: ToolCallFull, +} + +impl ToolcallStartPayload { + /// Creates a new tool call start payload + pub fn new(tool_call: ToolCallFull) -> Self { + Self { tool_call } + } +} + +/// Payload for the ToolcallEnd event +#[derive(Debug, PartialEq, Clone, Setters)] +#[setters(into)] +pub struct ToolcallEndPayload { + /// The original tool call that was executed + pub tool_call: ToolCallFull, + /// The tool result (success or failure) + pub result: ToolResult, +} + +impl ToolcallEndPayload { + /// Creates a new tool call end payload + pub fn new(tool_call: ToolCallFull, result: ToolResult) -> Self { + Self { tool_call, result } + } +} + +/// Lifecycle events that can occur during conversation processing +#[derive(Debug, PartialEq, Clone, From)] +pub enum LifecycleEvent { + /// Event fired when conversation processing starts + Start(EventData), + + /// Event fired when conversation processing ends + End(EventData), + + /// Event fired when a request is made to the LLM + Request(EventData), + + /// Event fired when a response is received from the LLM + Response(EventData), + + /// Event fired when a tool call starts + ToolcallStart(EventData), + + /// Event fired when a tool call ends + ToolcallEnd(EventData), +} + +/// Trait for handling lifecycle events +/// +/// Implementations of this trait can be used to react to different +/// stages of conversation processing. +#[async_trait] +pub trait EventHandle: Send + Sync { + /// Handles a lifecycle event and potentially modifies the conversation + /// + /// # Arguments + /// * `event` - The lifecycle event that occurred + /// * `conversation` - The current conversation state (mutable) + /// + /// # Errors + /// Returns an error if the event handling fails + async fn handle(&self, event: &T, conversation: &mut Conversation) -> anyhow::Result<()>; +} + +/// Extension trait for combining event handlers +/// +/// This trait provides methods to combine multiple event handlers into a single +/// handler that executes them in sequence. +pub trait EventHandleExt: EventHandle { + /// Combines this handler with another handler, creating a new handler that + /// runs both in sequence + /// + /// When an event is handled, both handlers run in sequence. + /// + /// # Arguments + /// * `other` - Another handler to combine with this one + /// + /// # Returns + /// A new boxed handler that combines both handlers + fn and + 'static>(self, other: H) -> Box> + where + Self: Sized + 'static; +} + +impl + 'static> EventHandleExt for A { + fn and + 'static>(self, other: H) -> Box> + where + Self: Sized + 'static, + { + Box::new(CombinedHandler(Box::new(self), Box::new(other))) + } +} + +// Implement EventHandle for Box to allow using boxed handlers +#[async_trait] +impl EventHandle for Box> { + async fn handle(&self, event: &T, conversation: &mut Conversation) -> anyhow::Result<()> { + (**self).handle(event, conversation).await + } +} + +/// A hook that contains handlers for all lifecycle events +/// +/// Hooks allow you to attach custom behavior at specific points +/// during conversation processing. +pub struct Hook { + on_start: Box>>, + on_end: Box>>, + on_request: Box>>, + on_response: Box>>, + on_toolcall_start: Box>>, + on_toolcall_end: Box>>, +} + +impl Default for Hook { + fn default() -> Self { + Self { + on_start: Box::new(NoOpHandler), + on_end: Box::new(NoOpHandler), + on_request: Box::new(NoOpHandler), + on_response: Box::new(NoOpHandler), + on_toolcall_start: Box::new(NoOpHandler), + on_toolcall_end: Box::new(NoOpHandler), + } + } +} + +impl Hook { + /// Creates a new hook with custom handlers for all event types + /// + /// # Arguments + /// * `on_start` - Handler for start events + /// * `on_end` - Handler for end events + /// * `on_request` - Handler for request events + /// * `on_response` - Handler for response events + /// * `on_toolcall_start` - Handler for tool call start events + /// * `on_toolcall_end` - Handler for tool call end events + pub fn new( + on_start: impl Into>>>, + on_end: impl Into>>>, + on_request: impl Into>>>, + on_response: impl Into>>>, + on_toolcall_start: impl Into>>>, + on_toolcall_end: impl Into>>>, + ) -> Self { + Self { + on_start: on_start.into(), + on_end: on_end.into(), + on_request: on_request.into(), + on_response: on_response.into(), + on_toolcall_start: on_toolcall_start.into(), + on_toolcall_end: on_toolcall_end.into(), + } + } +} + +impl Hook { + /// Sets the start event handler + /// + /// # Arguments + /// * `handler` - Handler for start events (automatically boxed) + pub fn on_start( + mut self, + handler: impl EventHandle> + 'static, + ) -> Self { + self.on_start = Box::new(handler); + self + } + + /// Sets the end event handler + /// + /// # Arguments + /// * `handler` - Handler for end events (automatically boxed) + pub fn on_end(mut self, handler: impl EventHandle> + 'static) -> Self { + self.on_end = Box::new(handler); + self + } + + /// Sets the request event handler + /// + /// # Arguments + /// * `handler` - Handler for request events (automatically boxed) + pub fn on_request( + mut self, + handler: impl EventHandle> + 'static, + ) -> Self { + self.on_request = Box::new(handler); + self + } + + /// Sets the response event handler + /// + /// # Arguments + /// * `handler` - Handler for response events (automatically boxed) + pub fn on_response( + mut self, + handler: impl EventHandle> + 'static, + ) -> Self { + self.on_response = Box::new(handler); + self + } + + /// Sets the tool call start event handler + /// + /// # Arguments + /// * `handler` - Handler for tool call start events (automatically boxed) + pub fn on_toolcall_start( + mut self, + handler: impl EventHandle> + 'static, + ) -> Self { + self.on_toolcall_start = Box::new(handler); + self + } + + /// Sets the tool call end event handler + /// + /// # Arguments + /// * `handler` - Handler for tool call end events (automatically boxed) + pub fn on_toolcall_end( + mut self, + handler: impl EventHandle> + 'static, + ) -> Self { + self.on_toolcall_end = Box::new(handler); + self + } +} + +impl Hook { + /// Combines this hook with another hook, creating a new hook that runs both + /// handlers in sequence + /// + /// When an event is handled, the first hook's handler runs first, then the + /// second hook's handler runs. + /// + /// # Arguments + /// * `other` - Another hook to combine with this one + /// + /// # Returns + /// A new hook that combines both hooks' handlers + pub fn zip(self, other: Hook) -> Self { + Self { + on_start: self.on_start.and(other.on_start), + on_end: self.on_end.and(other.on_end), + on_request: self.on_request.and(other.on_request), + on_response: self.on_response.and(other.on_response), + on_toolcall_start: self.on_toolcall_start.and(other.on_toolcall_start), + on_toolcall_end: self.on_toolcall_end.and(other.on_toolcall_end), + } + } +} + +// Implement EventHandle for Hook to allow hooks to handle LifecycleEvent +#[async_trait] +impl EventHandle for Hook { + async fn handle( + &self, + event: &LifecycleEvent, + conversation: &mut Conversation, + ) -> anyhow::Result<()> { + match &event { + LifecycleEvent::Start(data) => self.on_start.handle(data, conversation).await, + LifecycleEvent::End(data) => self.on_end.handle(data, conversation).await, + LifecycleEvent::Request(data) => self.on_request.handle(data, conversation).await, + LifecycleEvent::Response(data) => self.on_response.handle(data, conversation).await, + LifecycleEvent::ToolcallStart(data) => { + self.on_toolcall_start.handle(data, conversation).await + } + LifecycleEvent::ToolcallEnd(data) => { + self.on_toolcall_end.handle(data, conversation).await + } + } + } +} + +/// A handler that combines two event handlers with sequential execution +/// +/// Runs the first handler, then runs the second handler. +/// +/// This is used internally by the `Hook::zip` and `EventHandleExt::and` +/// methods. +struct CombinedHandler(Box>, Box>); + +#[async_trait] +impl EventHandle for CombinedHandler { + async fn handle(&self, event: &T, conversation: &mut Conversation) -> anyhow::Result<()> { + // Run the first handler + self.0.handle(event, conversation).await?; + // Run the second handler with the cloned event + self.1.handle(event, conversation).await + } +} + +/// A no-op handler that does nothing +/// +/// This is useful as a default handler when you only want to +/// handle specific events. +#[derive(Debug, Default)] +pub struct NoOpHandler; + +#[async_trait] +impl EventHandle for NoOpHandler { + async fn handle(&self, _: &T, _: &mut Conversation) -> anyhow::Result<()> { + Ok(()) + } +} + +#[async_trait] +impl EventHandle for F +where + F: Fn(&T, &mut Conversation) -> Fut + Send + Sync, + Fut: std::future::Future> + Send, +{ + async fn handle(&self, event: &T, conversation: &mut Conversation) -> anyhow::Result<()> { + (self)(event, conversation).await + } +} + +impl From for Box> +where + F: Fn(&T, &mut Conversation) -> Fut + Send + Sync + 'static, + Fut: std::future::Future> + Send + 'static, +{ + fn from(handler: F) -> Self { + Box::new(handler) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + use crate::{Agent, AgentId, Conversation, ModelId, ProviderId}; + + fn test_agent() -> Agent { + Agent::new( + AgentId::new("test_agent"), + ProviderId::FORGE, + ModelId::new("test-model"), + ) + } + + fn test_model_id() -> ModelId { + ModelId::new("test-model") + } + + #[test] + fn test_no_op_handler() { + let handler = NoOpHandler; + let conversation = Conversation::generate(); + + // This test just ensures NoOpHandler compiles and is constructible + let _ = handler; + let _ = conversation; + } + + #[tokio::test] + async fn test_hook_on_start() { + let events = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let events_clone = events.clone(); + + let hook = Hook::default().on_start( + move |event: &EventData, _conversation: &mut Conversation| { + let events = events_clone.clone(); + let event = event.clone(); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + }, + ); + + let mut conversation = Conversation::generate(); + + hook.handle( + &LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)), + &mut conversation, + ) + .await + .unwrap(); + + let handled = events.lock().unwrap(); + assert_eq!(handled.len(), 1); + assert_eq!( + handled[0], + EventData::new(test_agent(), test_model_id(), StartPayload) + ); + } + + #[tokio::test] + async fn test_hook_builder() { + let events = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + + let hook = Hook::default() + .on_start({ + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = LifecycleEvent::Start(event.clone()); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + } + }) + .on_end({ + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = LifecycleEvent::End(event.clone()); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + } + }) + .on_request({ + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = LifecycleEvent::Request(event.clone()); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + } + }); + + let mut conversation = Conversation::generate(); + + // Test Start event + hook.handle( + &LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)), + &mut conversation, + ) + .await + .unwrap(); + // Test End event + hook.handle( + &LifecycleEvent::End(EventData::new(test_agent(), test_model_id(), EndPayload)), + &mut conversation, + ) + .await + .unwrap(); + // Test Request event + hook.handle( + &LifecycleEvent::Request(EventData::new( + test_agent(), + test_model_id(), + RequestPayload::new(1), + )), + &mut conversation, + ) + .await + .unwrap(); + + let handled = events.lock().unwrap(); + assert_eq!(handled.len(), 3); + assert_eq!( + handled[0], + LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)) + ); + assert_eq!( + handled[1], + LifecycleEvent::End(EventData::new(test_agent(), test_model_id(), EndPayload)) + ); + assert_eq!( + handled[2], + LifecycleEvent::Request(EventData::new( + test_agent(), + test_model_id(), + RequestPayload::new(1) + )) + ); + } + + #[tokio::test] + async fn test_hook_all_events() { + let events = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + + let hook = Hook::new( + { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = LifecycleEvent::Start(event.clone()); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + } + }, + { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = LifecycleEvent::End(event.clone()); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + } + }, + { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = LifecycleEvent::Request(event.clone()); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + } + }, + { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = LifecycleEvent::Response(event.clone()); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + } + }, + { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = LifecycleEvent::ToolcallStart(event.clone()); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + } + }, + { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = LifecycleEvent::ToolcallEnd(event.clone()); + async move { + events.lock().unwrap().push(event); + Ok(()) + } + } + }, + ); + + let mut conversation = Conversation::generate(); + + let all_events = vec![ + LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)), + LifecycleEvent::End(EventData::new(test_agent(), test_model_id(), EndPayload)), + LifecycleEvent::Request(EventData::new( + test_agent(), + test_model_id(), + RequestPayload::new(1), + )), + LifecycleEvent::Response(EventData::new( + test_agent(), + test_model_id(), + ResponsePayload::new(ChatCompletionMessageFull { + content: "test".to_string(), + reasoning: None, + tool_calls: vec![], + thought_signature: None, + reasoning_details: None, + usage: crate::Usage::default(), + finish_reason: None, + phase: None, + }), + )), + LifecycleEvent::ToolcallStart(EventData::new( + test_agent(), + test_model_id(), + ToolcallStartPayload::new(ToolCallFull::new("test_tool")), + )), + LifecycleEvent::ToolcallEnd(EventData::new( + test_agent(), + test_model_id(), + ToolcallEndPayload::new( + ToolCallFull::new("test_tool"), + ToolResult::new("test_tool"), + ), + )), + ]; + + for event in all_events { + hook.handle(&event, &mut conversation).await.unwrap(); + } + + let handled = events.lock().unwrap(); + assert_eq!(handled.len(), 6); + } + + #[tokio::test] + async fn test_step_mutable_conversation() { + let title = std::sync::Arc::new(std::sync::Mutex::new(None)); + let hook = Hook::default().on_start({ + let title = title.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let title = title.clone(); + async move { + *title.lock().unwrap() = Some("Modified title".to_string()); + Ok(()) + } + } + }); + let mut conversation = Conversation::generate(); + + assert!(title.lock().unwrap().is_none()); + + hook.handle( + &LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)), + &mut conversation, + ) + .await + .unwrap(); + + assert_eq!(*title.lock().unwrap(), Some("Modified title".to_string())); + } + + #[test] + fn test_hook_default() { + let hook = Hook::default(); + + // Just ensure it compiles and is constructible + let _ = hook; + } + + #[tokio::test] + async fn test_hook_zip() { + let counter1 = std::sync::Arc::new(std::sync::Mutex::new(0)); + let counter2 = std::sync::Arc::new(std::sync::Mutex::new(0)); + + let hook1 = Hook::default().on_start({ + let counter = counter1.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let counter = counter.clone(); + async move { + *counter.lock().unwrap() += 1; + Ok(()) + } + } + }); + + let hook2 = Hook::default().on_start({ + let counter = counter2.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let counter = counter.clone(); + async move { + *counter.lock().unwrap() += 1; + Ok(()) + } + } + }); + let combined: Hook = hook1.zip(hook2); + + let mut conversation = Conversation::generate(); + combined + .handle( + &LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)), + &mut conversation, + ) + .await + .unwrap(); + + // Both handlers should have been called + assert_eq!(*counter1.lock().unwrap(), 1); + assert_eq!(*counter2.lock().unwrap(), 1); + } + + #[tokio::test] + async fn test_hook_zip_multiple() { + let events = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + + let hook1 = Hook::default().on_start({ + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = event.clone(); + async move { + events.lock().unwrap().push(format!("h1:{:?}", event)); + Ok(()) + } + } + }); + + let hook2 = Hook::default().on_start({ + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = event.clone(); + async move { + events.lock().unwrap().push(format!("h2:{:?}", event)); + Ok(()) + } + } + }); + + let hook3 = Hook::default().on_start({ + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = event.clone(); + async move { + events.lock().unwrap().push(format!("h3:{:?}", event)); + Ok(()) + } + } + }); + let combined: Hook = hook1.zip(hook2).zip(hook3); + + let mut conversation = Conversation::generate(); + combined + .handle( + &LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)), + &mut conversation, + ) + .await + .unwrap(); + + let handled = events.lock().unwrap(); + assert_eq!(handled.len(), 3); + assert!(handled[0].starts_with("h1:EventData")); + assert!(handled[1].starts_with("h2:EventData")); + assert!(handled[2].starts_with("h3:EventData")); + } + + #[tokio::test] + async fn test_hook_zip_different_events() { + let start_title = std::sync::Arc::new(std::sync::Mutex::new(None)); + let end_title = std::sync::Arc::new(std::sync::Mutex::new(None)); + + let hook1 = Hook::default() + .on_start({ + let start_title = start_title.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let start_title = start_title.clone(); + async move { + *start_title.lock().unwrap() = Some("Start".to_string()); + Ok(()) + } + } + }) + .on_end({ + let end_title = end_title.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let end_title = end_title.clone(); + async move { + *end_title.lock().unwrap() = Some("End".to_string()); + Ok(()) + } + } + }); + let hook2 = Hook::default(); + + let combined: Hook = hook1.zip(hook2); + + let mut conversation = Conversation::generate(); + + // Test Start event + combined + .handle( + &LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)), + &mut conversation, + ) + .await + .unwrap(); + assert_eq!(*start_title.lock().unwrap(), Some("Start".to_string())); + + // Test End event + combined + .handle( + &LifecycleEvent::End(EventData::new(test_agent(), test_model_id(), EndPayload)), + &mut conversation, + ) + .await + .unwrap(); + assert_eq!(*end_title.lock().unwrap(), Some("End".to_string())); + } + + #[tokio::test] + async fn test_event_handle_ext_and() { + let counter1 = std::sync::Arc::new(std::sync::Mutex::new(0)); + let counter2 = std::sync::Arc::new(std::sync::Mutex::new(0)); + + let handler1 = { + let counter = counter1.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let counter = counter.clone(); + async move { + *counter.lock().unwrap() += 1; + Ok(()) + } + } + }; + + let handler2 = { + let counter = counter2.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let counter = counter.clone(); + async move { + *counter.lock().unwrap() += 1; + Ok(()) + } + } + }; + + let combined: Box>> = handler1.and(handler2); + + let mut conversation = Conversation::generate(); + combined + .handle( + &EventData::new(test_agent(), test_model_id(), StartPayload), + &mut conversation, + ) + .await + .unwrap(); + + // Both handlers should have been called + assert_eq!(*counter1.lock().unwrap(), 1); + assert_eq!(*counter2.lock().unwrap(), 1); + } + + #[tokio::test] + async fn test_event_handle_ext_and_boxed() { + let counter1 = std::sync::Arc::new(std::sync::Mutex::new(0)); + let counter2 = std::sync::Arc::new(std::sync::Mutex::new(0)); + + let handler1 = { + let counter = counter1.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let counter = counter.clone(); + async move { + *counter.lock().unwrap() += 1; + Ok(()) + } + } + }; + + let handler2 = { + let counter = counter2.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let counter = counter.clone(); + async move { + *counter.lock().unwrap() += 1; + Ok(()) + } + } + }; + + let combined: Box>> = handler1.and(handler2); + + let mut conversation = Conversation::generate(); + combined + .handle( + &EventData::new(test_agent(), test_model_id(), StartPayload), + &mut conversation, + ) + .await + .unwrap(); + + // Both handlers should have been called + assert_eq!(*counter1.lock().unwrap(), 1); + assert_eq!(*counter2.lock().unwrap(), 1); + } + + #[tokio::test] + async fn test_event_handle_ext_chain() { + let events = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + + let handler1 = { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = event.clone(); + async move { + events.lock().unwrap().push(format!("h1:{:?}", event)); + Ok(()) + } + } + }; + + let handler2 = { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = event.clone(); + async move { + events.lock().unwrap().push(format!("h2:{:?}", event)); + Ok(()) + } + } + }; + + let handler3 = { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = event.clone(); + async move { + events.lock().unwrap().push(format!("h3:{:?}", event)); + Ok(()) + } + } + }; + + // Chain handlers using and() + let combined: Box>> = + handler1.and(handler2).and(handler3); + + let mut conversation = Conversation::generate(); + combined + .handle( + &EventData::new(test_agent(), test_model_id(), StartPayload), + &mut conversation, + ) + .await + .unwrap(); + + let handled = events.lock().unwrap(); + assert_eq!(handled.len(), 3); + assert!(handled[0].starts_with("h1:EventData")); + assert!(handled[1].starts_with("h2:EventData")); + assert!(handled[2].starts_with("h3:EventData")); + } + + #[tokio::test] + async fn test_event_handle_ext_with_hook() { + let events = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let start_title = std::sync::Arc::new(std::sync::Mutex::new(None)); + + let start_handler = { + let start_title = start_title.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let start_title = start_title.clone(); + async move { + *start_title.lock().unwrap() = Some("Started".to_string()); + Ok(()) + } + } + }; + + let logging_handler = { + let events = events.clone(); + move |event: &EventData, _conversation: &mut Conversation| { + let events = events.clone(); + let event = event.clone(); + async move { + events.lock().unwrap().push(format!("Event: {:?}", event)); + Ok(()) + } + } + }; + + // Combine handlers using extension trait + let combined_handler: Box>> = + start_handler.and(logging_handler); + + let hook = Hook::default().on_start(combined_handler); + + let mut conversation = Conversation::generate(); + hook.handle( + &LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)), + &mut conversation, + ) + .await + .unwrap(); + + assert_eq!(events.lock().unwrap().len(), 1); + assert!(events.lock().unwrap()[0].starts_with("Event: EventData")); + } + + #[tokio::test] + async fn test_hook_as_event_handle() { + let start_title = std::sync::Arc::new(std::sync::Mutex::new(None)); + let end_title = std::sync::Arc::new(std::sync::Mutex::new(None)); + + let hook = Hook::default() + .on_start({ + let start_title = start_title.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let start_title = start_title.clone(); + async move { + *start_title.lock().unwrap() = Some("Started".to_string()); + Ok(()) + } + } + }) + .on_end({ + let end_title = end_title.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let end_title = end_title.clone(); + async move { + *end_title.lock().unwrap() = Some("Ended".to_string()); + Ok(()) + } + } + }); + + // Test using handle() directly (EventHandle trait) + let mut conversation = Conversation::generate(); + hook.handle( + &LifecycleEvent::Start(EventData::new(test_agent(), test_model_id(), StartPayload)), + &mut conversation, + ) + .await + .unwrap(); + assert_eq!(*start_title.lock().unwrap(), Some("Started".to_string())); + + hook.handle( + &LifecycleEvent::End(EventData::new(test_agent(), test_model_id(), EndPayload)), + &mut conversation, + ) + .await + .unwrap(); + assert_eq!(*end_title.lock().unwrap(), Some("Ended".to_string())); + } + + #[tokio::test] + async fn test_hook_combination_with_and() { + let hook1_title = std::sync::Arc::new(std::sync::Mutex::new(None)); + let hook2_title = std::sync::Arc::new(std::sync::Mutex::new(None)); + + let handler1 = { + let hook1_title = hook1_title.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let hook1_title = hook1_title.clone(); + async move { + *hook1_title.lock().unwrap() = Some("Started".to_string()); + Ok(()) + } + } + }; + let handler2 = { + let hook2_title = hook2_title.clone(); + move |_event: &EventData, _conversation: &mut Conversation| { + let hook2_title = hook2_title.clone(); + async move { + *hook2_title.lock().unwrap() = Some("Ended".to_string()); + Ok(()) + } + } + }; + + // Combine handlers using and() extension method + let combined: Box>> = handler1.and(handler2); + + let mut conversation = Conversation::generate(); + combined + .handle( + &EventData::new(test_agent(), test_model_id(), StartPayload), + &mut conversation, + ) + .await + .unwrap(); + + // Both handlers should have been called + assert_eq!(*hook1_title.lock().unwrap(), Some("Started".to_string())); + assert_eq!(*hook2_title.lock().unwrap(), Some("Ended".to_string())); + } +} diff --git a/crates/forge_domain/src/http_config.rs b/crates/forge_domain/src/http_config.rs new file mode 100644 index 0000000000000000000000000000000000000000..17ac53c702838b34ecb1bc89415c3eedb5dcfc2e --- /dev/null +++ b/crates/forge_domain/src/http_config.rs @@ -0,0 +1,284 @@ +use serde::{Deserialize, Serialize}; +use strum_macros::EnumString; + +/// TLS version enum for configuring minimum and maximum TLS protocol versions. +/// +/// Used with `HttpConfig` to specify TLS version constraints for HTTP +/// connections. +/// +/// # Example +/// ``` +/// use forge_domain::{HttpConfig, TlsVersion, TlsBackend}; +/// +/// let config = HttpConfig { +/// min_tls_version: Some(TlsVersion::V1_2), +/// max_tls_version: Some(TlsVersion::V1_3), +/// tls_backend: TlsBackend::Rustls, +/// ..HttpConfig::default() +/// }; +/// ``` +/// +/// # Environment Variables +/// - `FORGE_HTTP_MIN_TLS_VERSION`: Set minimum TLS version (e.g., "1.2") +/// - `FORGE_HTTP_MAX_TLS_VERSION`: Set maximum TLS version (e.g., "1.3") +#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize, fake::Dummy)] +#[serde(rename_all = "camelCase")] +pub enum TlsVersion { + #[serde(rename = "1.0")] + V1_0, + #[serde(rename = "1.1")] + V1_1, + #[serde(rename = "1.2")] + V1_2, + #[default] + #[serde(rename = "1.3")] + V1_3, +} + +impl std::fmt::Display for TlsVersion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TlsVersion::V1_0 => write!(f, "1.0"), + TlsVersion::V1_1 => write!(f, "1.1"), + TlsVersion::V1_2 => write!(f, "1.2"), + TlsVersion::V1_3 => write!(f, "1.3"), + } + } +} + +impl std::str::FromStr for TlsVersion { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + "1.0" => Ok(TlsVersion::V1_0), + "1.1" => Ok(TlsVersion::V1_1), + "1.2" => Ok(TlsVersion::V1_2), + "1.3" => Ok(TlsVersion::V1_3), + _ => Err(format!( + "Invalid TLS version: {s}. Valid options are: 1.0, 1.1, 1.2, 1.3" + )), + } + } +} + +#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize, EnumString, fake::Dummy)] +#[serde(rename_all = "camelCase")] +#[strum(serialize_all = "lowercase")] +pub enum TlsBackend { + #[default] + Default, + Rustls, +} + +impl std::fmt::Display for TlsBackend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TlsBackend::Default => write!(f, "default"), + TlsBackend::Rustls => write!(f, "rustls"), + } + } +} + +/// HTTP client configuration with support for timeouts, connection pooling, +/// redirects, DNS resolution, TLS settings, and HTTP/2 configuration. +/// +/// # TLS Configuration +/// The `min_tls_version` and `max_tls_version` fields allow you to specify +/// TLS protocol version constraints. These are optional and when `None`, +/// the TLS library defaults will be used. +/// +/// # HTTP/2 Configuration +/// The HTTP/2 settings control adaptive window sizing, keep-alive behavior, +/// and connection management for HTTP/2 connections. +/// +/// # Environment Variables +/// All HttpConfig fields can be configured via environment variables: +/// - `FORGE_HTTP_CONNECT_TIMEOUT`: Connection timeout in seconds (default: 30) +/// - `FORGE_HTTP_READ_TIMEOUT`: Read timeout in seconds (default: 900) +/// - `FORGE_HTTP_POOL_IDLE_TIMEOUT`: Pool idle timeout in seconds (default: 90) +/// - `FORGE_HTTP_POOL_MAX_IDLE_PER_HOST`: Max idle connections per host +/// (default: 5) +/// - `FORGE_HTTP_MAX_REDIRECTS`: Maximum redirects to follow (default: 10) +/// - `FORGE_HTTP_USE_HICKORY`: Use Hickory DNS resolver (default: false) +/// - `FORGE_HTTP_TLS_BACKEND`: TLS backend ("default" or "rustls", default: +/// "default") +/// - `FORGE_HTTP_MIN_TLS_VERSION`: Minimum TLS version ("1.0", "1.1", "1.2", +/// "1.3") +/// - `FORGE_HTTP_MAX_TLS_VERSION`: Maximum TLS version ("1.0", "1.1", "1.2", +/// "1.3") +/// - `FORGE_HTTP_ADAPTIVE_WINDOW`: Enable HTTP/2 adaptive window (default: +/// true) +/// - `FORGE_HTTP_KEEP_ALIVE_INTERVAL`: Keep-alive interval in seconds (default: +/// 60, use "none"/"disabled" to disable) +/// - `FORGE_HTTP_KEEP_ALIVE_TIMEOUT`: Keep-alive timeout in seconds (default: +/// 10) +/// - `FORGE_HTTP_KEEP_ALIVE_WHILE_IDLE`: Keep-alive while idle (default: true) +/// - `FORGE_HTTP_ACCEPT_INVALID_CERTS`: Accept invalid certificates (default: +/// false) - USE WITH CAUTION +/// - `FORGE_HTTP_ROOT_CERT_PATHS`: Paths to root certificate files (PEM, CRT, +/// CER format), multiple paths separated by commas +/// +/// # Example +/// ``` +/// use forge_domain::{HttpConfig, TlsVersion, TlsBackend}; +/// +/// let config = HttpConfig { +/// connect_timeout: 30, +/// min_tls_version: Some(TlsVersion::V1_2), +/// max_tls_version: Some(TlsVersion::V1_3), +/// tls_backend: TlsBackend::Rustls, +/// adaptive_window: true, +/// keep_alive_interval: Some(60), +/// ..HttpConfig::default() +/// }; +/// ``` +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, fake::Dummy)] +#[serde(rename_all = "camelCase")] +pub struct HttpConfig { + pub connect_timeout: u64, + pub read_timeout: u64, + pub pool_idle_timeout: u64, + pub pool_max_idle_per_host: usize, + pub max_redirects: usize, + pub hickory: bool, + pub tls_backend: TlsBackend, + /// Minimum TLS protocol version to use. When `None`, uses TLS library + /// default. + pub min_tls_version: Option, + /// Maximum TLS protocol version to use. When `None`, uses TLS library + /// default. + pub max_tls_version: Option, + /// Adaptive window sizing for improved flow control. + pub adaptive_window: bool, + /// Keep-alive interval in seconds. When `None`, keep-alive is + /// disabled. + pub keep_alive_interval: Option, + /// Keep-alive timeout in seconds. + pub keep_alive_timeout: u64, + /// Keep-alive while connection is idle. + pub keep_alive_while_idle: bool, + /// Accept invalid certificates. This should be used with caution. + pub accept_invalid_certs: bool, + /// Paths to root certificate files (PEM, CRT, CER format). Multiple paths + /// can be separated by commas. + pub root_cert_paths: Option>, +} + +impl Default for HttpConfig { + fn default() -> Self { + Self { + connect_timeout: 30, // 30 seconds + read_timeout: 900, /* 15 minutes; this should be in sync with the server function + * execution timeout */ + pool_idle_timeout: 90, + pool_max_idle_per_host: 5, + max_redirects: 10, + hickory: false, + tls_backend: TlsBackend::default(), + min_tls_version: None, // Use TLS library default + max_tls_version: None, // Use TLS library default + // HTTP/2 defaults - enable HTTP/2 with sensible keep-alive settings + adaptive_window: true, + keep_alive_interval: Some(60), // 60 seconds + keep_alive_timeout: 10, // 10 seconds + keep_alive_while_idle: true, + accept_invalid_certs: false, // Default to false for security + root_cert_paths: None, + } + } +} + +#[cfg(test)] +mod tests { + use std::str::FromStr; + + use super::*; + + #[test] + fn test_tls_version_from_str() { + assert_eq!(TlsVersion::from_str("1.0").unwrap(), TlsVersion::V1_0); + assert_eq!(TlsVersion::from_str("1.1").unwrap(), TlsVersion::V1_1); + assert_eq!(TlsVersion::from_str("1.2").unwrap(), TlsVersion::V1_2); + assert_eq!(TlsVersion::from_str("1.3").unwrap(), TlsVersion::V1_3); + + assert!(TlsVersion::from_str("invalid").is_err()); + assert!(TlsVersion::from_str("2.0").is_err()); + } + + #[test] + fn test_tls_version_display() { + assert_eq!(TlsVersion::V1_0.to_string(), "1.0"); + assert_eq!(TlsVersion::V1_1.to_string(), "1.1"); + assert_eq!(TlsVersion::V1_2.to_string(), "1.2"); + assert_eq!(TlsVersion::V1_3.to_string(), "1.3"); + } + + #[test] + fn test_tls_version_default() { + assert_eq!(TlsVersion::default(), TlsVersion::V1_3); + } + + #[test] + fn test_http_config_with_tls_versions() { + let config = HttpConfig { + min_tls_version: Some(TlsVersion::V1_2), + max_tls_version: Some(TlsVersion::V1_3), + ..HttpConfig::default() + }; + + assert_eq!(config.min_tls_version, Some(TlsVersion::V1_2)); + assert_eq!(config.max_tls_version, Some(TlsVersion::V1_3)); + } + + #[test] + fn test_http_config_http2_defaults() { + let config = HttpConfig::default(); + + assert!(config.adaptive_window); + assert_eq!(config.keep_alive_interval, Some(60)); + assert_eq!(config.keep_alive_timeout, 10); + assert!(config.keep_alive_while_idle); + } + + #[test] + fn test_http_config_http2_custom_values() { + let config = HttpConfig { + adaptive_window: false, + keep_alive_interval: None, + keep_alive_timeout: 30, + keep_alive_while_idle: false, + ..HttpConfig::default() + }; + + assert!(!config.adaptive_window); + assert_eq!(config.keep_alive_interval, None); + assert_eq!(config.keep_alive_timeout, 30); + assert!(!config.keep_alive_while_idle); + } + + #[test] + fn test_http_config_accept_invalid_certs_defaults() { + let config = HttpConfig::default(); + assert!(!config.accept_invalid_certs); + } + + #[test] + fn test_http_config_accept_invalid_certs_custom() { + let config = HttpConfig { accept_invalid_certs: true, ..HttpConfig::default() }; + assert!(config.accept_invalid_certs); + } + + #[test] + fn test_http_config_root_cert_paths_custom() { + let cert_paths = vec![ + "/path/to/cert1.pem".to_string(), + "/path/to/cert2.crt".to_string(), + ]; + let config = HttpConfig { + root_cert_paths: Some(cert_paths.clone()), + ..HttpConfig::default() + }; + assert_eq!(config.root_cert_paths, Some(cert_paths)); + } +} diff --git a/crates/forge_domain/src/image.rs b/crates/forge_domain/src/image.rs new file mode 100644 index 0000000000000000000000000000000000000000..a4770e4606789da9e09eeed2e23a635bcb93a4cf --- /dev/null +++ b/crates/forge_domain/src/image.rs @@ -0,0 +1,30 @@ +use base64::Engine; +use derive_getters::Getters; +use serde::{Deserialize, Serialize}; + +#[derive(Default, Clone, Debug, Serialize, Deserialize, Getters, PartialEq, Eq, Hash)] +pub struct Image { + url: String, + mime_type: String, +} + +impl Image { + pub fn new_bytes(content: Vec, mime_type: impl ToString) -> Self { + let mime_type = mime_type.to_string(); + let base64_encoded = base64::engine::general_purpose::STANDARD.encode(&content); + Self::new_base64(base64_encoded, mime_type) + } + + // Returns the base64 image without the prefix. + pub fn data(&self) -> &str { + self.url + .strip_prefix(&format!("data:{};base64,", self.mime_type)) + .unwrap_or(&self.url) + } + + pub fn new_base64(base64_encoded: String, mime_type: impl ToString) -> Self { + let mime_type = mime_type.to_string(); + let content = format!("data:{mime_type};base64,{base64_encoded}"); + Self { url: content, mime_type } + } +} diff --git a/crates/forge_domain/src/lib.rs b/crates/forge_domain/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..5ae3fca85dd42ae3daba565cc4a602b6a8b8ef56 --- /dev/null +++ b/crates/forge_domain/src/lib.rs @@ -0,0 +1,119 @@ +mod agent; +mod attachment; +mod auth; +mod chat_request; +mod chat_response; +mod command; +mod compact; +mod console; +mod context; +mod conversation; +mod conversation_html; +mod data_gen; +mod env; +mod error; +mod event; +mod file; +mod file_operation; +mod group_by_key; +mod hook; +mod http_config; +mod image; +mod max_tokens; +mod mcp; +mod mcp_servers; +mod merge; +mod message; +mod message_pattern; +mod migration; +mod model; +mod model_config; +mod node; +mod point; +mod policies; +mod provider; +mod reasoning; +mod repo; +mod result_stream_ext; +mod session_metrics; +mod shell; +mod skill; +mod snapshot; +mod suggestion; +mod system_context; +mod temperature; +mod template; +mod terminal_context; +mod tools; + +mod tool_order; +mod top_k; +mod top_p; +mod transformer; +mod update; +mod validation; +mod workspace; +mod xml; + +pub use agent::*; +pub use attachment::*; +pub use chat_request::*; +pub use chat_response::*; +pub use command::*; +pub use compact::*; +pub use console::*; +pub use context::*; +pub use conversation::*; +pub use conversation_html::*; +pub use data_gen::*; +pub use env::*; +pub use error::*; +pub use event::*; +pub use file::*; +pub use file_operation::*; +pub use fuzzy_search::*; +pub use group_by_key::*; +pub use hook::*; +pub use http_config::*; +pub use image::*; +pub use max_tokens::*; +pub use mcp::*; +pub use mcp_servers::*; +pub use message::*; +pub use message_pattern::*; +pub use migration::*; +pub use model::*; +pub use model_config::*; +pub use node::*; +pub use point::*; +pub use policies::*; +pub use provider::*; +pub use reasoning::*; +pub use repo::*; +pub use result_stream_ext::*; +pub use session_metrics::*; +pub use shell::*; +pub use skill::*; +pub use snapshot::*; +pub use suggestion::*; +pub use system_context::*; +pub use temperature::*; +pub use template::*; +pub use terminal_context::*; +pub use tool_order::*; +pub use tools::*; +pub use top_k::*; +pub use top_p::*; +pub use transformer::*; +pub use update::*; +pub use validation::*; +pub use workspace::*; +pub use xml::*; +pub mod line_numbers; +pub use auth::*; +pub use line_numbers::*; + +/// Type alias for tokio mpsc Sender that sends ChatResponse results +pub type ArcSender = tokio::sync::mpsc::Sender>; + +mod fuzzy_search; diff --git a/crates/forge_domain/src/max_tokens.rs b/crates/forge_domain/src/max_tokens.rs new file mode 100644 index 0000000000000000000000000000000000000000..b93d436e200e9e9b23e56dfc0544367bf926d053 --- /dev/null +++ b/crates/forge_domain/src/max_tokens.rs @@ -0,0 +1,199 @@ +use std::fmt; +use std::ops::Deref; + +use schemars::JsonSchema; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +/// A newtype for max_tokens values with built-in validation +/// +/// Max tokens controls the maximum number of tokens the model can generate: +/// - Lower values (e.g., 100) limit response length for concise outputs +/// - Higher values (e.g., 4000) allow for longer, more detailed responses +/// - Valid range is 1 to 100,000 (reasonable upper bound for most models) +/// - If not specified, the model provider's default will be used +#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, JsonSchema)] +pub struct MaxTokens(u32); + +impl MaxTokens { + /// Creates a new MaxTokens value, returning an error if outside the valid + /// range (1 to 100,000) + pub fn new(value: u32) -> Result { + if Self::is_valid(value) { + Ok(Self(value)) + } else { + Err(format!( + "max_tokens must be between 1 and 100000, got {value}" + )) + } + } + + /// Creates a new MaxTokens value without validation + /// + /// # Safety + /// This function should only be used when the value is known to be valid + pub fn new_unchecked(value: u32) -> Self { + debug_assert!(Self::is_valid(value), "invalid max_tokens: {value}"); + Self(value) + } + + /// Returns true if the max_tokens value is within the valid range (1 to + /// 100,000) + pub fn is_valid(value: u32) -> bool { + (1..=100_000).contains(&value) + } + + /// Returns the inner u32 value + pub fn value(&self) -> u32 { + self.0 + } +} + +impl Deref for MaxTokens { + type Target = u32; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl From for u32 { + fn from(max_tokens: MaxTokens) -> Self { + max_tokens.0 + } +} + +impl fmt::Display for MaxTokens { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +impl Serialize for MaxTokens { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_u32(self.0) + } +} + +impl<'de> Deserialize<'de> for MaxTokens { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + use serde::de::Error; + let value = u32::deserialize(deserializer)?; + if Self::is_valid(value) { + Ok(Self(value)) + } else { + Err(Error::custom(format!( + "max_tokens must be between 1 and 100000, got {value}" + ))) + } + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use serde_json::json; + + use super::*; + + #[test] + fn test_max_tokens_creation() { + // Valid max_tokens values should be created successfully + let valid_values = [1, 100, 1000, 4000, 8000, 100_000]; + for value in valid_values { + let result = MaxTokens::new(value); + assert!(result.is_ok(), "MaxTokens {value} should be valid"); + assert_eq!(result.unwrap().value(), value); + } + + // Invalid max_tokens values should return an error + let invalid_values = [0, 100_001, 200_000, 1_000_000]; + for value in invalid_values { + let result = MaxTokens::new(value); + assert!(result.is_err(), "MaxTokens {value} should be invalid"); + assert!( + result + .unwrap_err() + .contains("max_tokens must be between 1 and 100000"), + "Error should mention valid range" + ); + } + } + + #[test] + fn test_max_tokens_serialization() { + let max_tokens = MaxTokens::new(4000).unwrap(); + let json = serde_json::to_value(max_tokens).unwrap(); + + if let serde_json::Value::Number(num) = &json { + let int_val = num.as_u64().unwrap(); + assert_eq!(int_val, 4000); + } else { + panic!("Expected a number, got {json:?}"); + } + } + + #[test] + fn test_max_tokens_deserialization() { + // Valid max_tokens values should deserialize correctly + let valid_values = [1, 100, 1000, 4000, 8000, 100_000]; + for value in valid_values { + let json = json!(value); + let max_tokens: Result = serde_json::from_value(json); + assert!( + max_tokens.is_ok(), + "Valid max_tokens {value} should deserialize" + ); + assert_eq!(max_tokens.unwrap().value(), value); + } + + // Invalid max_tokens values should fail deserialization + let invalid_values = [0, 100_001, 200_000, 1_000_000]; + for value in invalid_values { + let json = json!(value); + let max_tokens: Result = serde_json::from_value(json); + assert!( + max_tokens.is_err(), + "Invalid max_tokens {value} should fail deserialization" + ); + let err = max_tokens.unwrap_err().to_string(); + assert!( + err.contains("max_tokens must be between 1 and 100000"), + "Error should mention valid range: {err}" + ); + } + } + + #[test] + fn test_max_tokens_in_struct() { + #[derive(Serialize, Deserialize, Debug)] + struct TestStruct { + max_tokens: MaxTokens, + } + + // Valid max_tokens + let json = json!({ + "max_tokens": 4000 + }); + let test_struct: Result = serde_json::from_value(json); + assert!(test_struct.is_ok()); + assert_eq!(test_struct.unwrap().max_tokens.value(), 4000); + + // Invalid max_tokens + let json = json!({ + "max_tokens": 200_000 + }); + let test_struct: Result = serde_json::from_value(json); + assert!(test_struct.is_err()); + let err = test_struct.unwrap_err().to_string(); + assert!( + err.contains("max_tokens must be between 1 and 100000"), + "Error should mention valid range: {err}" + ); + } +} diff --git a/crates/forge_domain/src/mcp_servers.rs b/crates/forge_domain/src/mcp_servers.rs new file mode 100644 index 0000000000000000000000000000000000000000..68d2b1ae8b652e414c58f8905855afa62fbdd6ac --- /dev/null +++ b/crates/forge_domain/src/mcp_servers.rs @@ -0,0 +1,50 @@ +use std::collections::{HashMap, hash_map}; + +use serde::{Deserialize, Serialize}; + +use crate::{ServerName, ToolDefinition}; + +/// Cache for MCP tool definitions +/// +/// Simplified cache structure that stores only the essential data. +/// Validation and TTL checking are handled by the infrastructure layer +/// using cacache's built-in metadata capabilities. +#[derive(Default, Clone, Serialize, Deserialize, Debug, PartialEq, derive_setters::Setters)] +#[serde(rename_all = "camelCase")] +#[setters(strip_option, into)] +pub struct McpServers { + /// Successfully loaded MCP servers with their tools + servers: HashMap>, + /// Failed MCP servers with their error messages + #[serde(default)] + failures: HashMap, +} + +impl McpServers { + /// Create a new cache entry with servers and failures + pub fn new( + servers: HashMap>, + failures: HashMap, + ) -> Self { + Self { servers, failures } + } + + /// Get the successful servers + pub fn get_servers(&self) -> &HashMap> { + &self.servers + } + + /// Get the failed servers + pub fn get_failures(&self) -> &HashMap { + &self.failures + } +} + +impl IntoIterator for McpServers { + type Item = (ServerName, Vec); + type IntoIter = hash_map::IntoIter>; + + fn into_iter(self) -> Self::IntoIter { + self.servers.into_iter() + } +} diff --git a/crates/forge_domain/src/merge.rs b/crates/forge_domain/src/merge.rs new file mode 100644 index 0000000000000000000000000000000000000000..faea53202e4e279f606ae1712e980904bcdd71ae --- /dev/null +++ b/crates/forge_domain/src/merge.rs @@ -0,0 +1,59 @@ +use ::std::collections::HashMap; +use ::std::hash::Hash; + +pub mod std { + pub fn overwrite(base: &mut T, other: T) { + *base = other; + } +} + +pub mod vec { + + use std::collections::HashMap; + + use merge::Merge; + + use super::Key; + + #[allow(unused)] + pub fn unify_by_key(base: &mut Vec, other: Vec) + where + T::Id: Clone + std::hash::Hash + Eq, + { + // Create a HashMap for O(1) lookup of base agents by their key + let mut base_map: HashMap = HashMap::new(); + for (index, agent) in base.iter().enumerate() { + base_map.insert(agent.key().clone(), index); + } + + for other_agent in other { + if let Some(&index) = base_map.get(other_agent.key()) { + // If the base contains an agent with the same Key, merge them + if let Some(base_agent) = base.get_mut(index) { + base_agent.merge(other_agent); + } + } else { + // Otherwise, append the other agent to the base list + base.push(other_agent); + } + } + } +} + +pub fn option(base: &mut Option, other: Option) { + if other.is_some() { + *base = other; + } +} +#[allow(unused)] +pub trait Key { + type Id: Eq; + fn key(&self) -> &Self::Id; +} + +#[allow(unused)] +pub fn hashmap(base: &mut HashMap, other: HashMap) { + for (key, value) in other { + base.insert(key, value); + } +} diff --git a/crates/forge_domain/src/message.rs b/crates/forge_domain/src/message.rs new file mode 100644 index 0000000000000000000000000000000000000000..38440ef0610c1b09afcbab36a6003f0b4cb0c47e --- /dev/null +++ b/crates/forge_domain/src/message.rs @@ -0,0 +1,471 @@ +use derive_more::derive::From; +use derive_setters::Setters; +use serde::{Deserialize, Serialize}; +use strum_macros::{EnumString, IntoStaticStr}; + +use super::{ToolCall, ToolCallFull}; +use crate::TokenCount; +use crate::reasoning::{Reasoning, ReasoningFull}; + +/// Labels an assistant message as intermediate commentary or the final answer. +/// +/// For models like `gpt-5.3-codex` and beyond, when sending follow-up requests, +/// preserve and resend phase on all assistant messages -- dropping it can +/// degrade performance. +#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum MessagePhase { + /// Intermediate commentary produced while the model is reasoning. + Commentary, + /// The final answer from the model. + FinalAnswer, +} + +#[derive(Default, Clone, Copy, Debug, Serialize, Deserialize, PartialEq)] +pub struct Usage { + pub prompt_tokens: TokenCount, + pub completion_tokens: TokenCount, + pub total_tokens: TokenCount, + pub cached_tokens: TokenCount, + pub cost: Option, +} + +impl Usage { + /// Accumulates usage from another Usage instance by summing all fields. + /// + /// Use this for aggregating usage across **independent** requests (e.g., + /// session-level totals where each message has its own final usage). + pub fn accumulate(mut self, other: &Usage) -> Self { + self.prompt_tokens = self.prompt_tokens + other.prompt_tokens; + self.completion_tokens = self.completion_tokens + other.completion_tokens; + self.total_tokens = self.total_tokens + other.total_tokens; + self.cached_tokens = self.cached_tokens + other.cached_tokens; + self.cost = match (self.cost, other.cost) { + (Some(a), Some(b)) => Some(a + b), + (Some(a), None) => Some(a), + (None, Some(b)) => Some(b), + (None, None) => None, + }; + self + } + + /// Merges usage from another Usage instance using a "last non-zero wins" + /// strategy. + /// + /// Use this when combining **partial** usage events within a single + /// streaming response where values are **cumulative** (not incremental): + /// - `message_start`: `input_tokens=1000, output_tokens=1` + /// - `message_delta`: `input_tokens=0, output_tokens=75` (cumulative + /// total) + /// + /// For each field, the larger of the two values is kept. This prevents + /// double-counting when providers report cumulative token counts across + /// multiple events. + /// + /// Cost is summed since cost events are always additive. + pub fn merge(mut self, other: &Usage) -> Self { + self.prompt_tokens = self.prompt_tokens.max(other.prompt_tokens); + self.completion_tokens = self.completion_tokens.max(other.completion_tokens); + self.total_tokens = self.total_tokens.max(other.total_tokens); + self.cached_tokens = self.cached_tokens.max(other.cached_tokens); + self.cost = match (self.cost, other.cost) { + (Some(a), Some(b)) => Some(a + b), + (Some(a), None) => Some(a), + (None, Some(b)) => Some(b), + (None, None) => None, + }; + self + } +} + +/// Represents a message that was received from the LLM provider +/// NOTE: Tool call messages are part of the larger Response object and not part +/// of the message. +#[derive(Default, Clone, Debug, Setters, PartialEq)] +#[setters(into, strip_option)] +pub struct ChatCompletionMessage { + pub content: Option, + pub thought_signature: Option, + pub reasoning: Option, + pub reasoning_details: Option>, + pub tool_calls: Vec, + pub finish_reason: Option, + pub usage: Option, + /// Phase label for assistant messages (e.g. `Commentary` or `FinalAnswer`). + /// Preserved from the response and replayed back on subsequent requests. + pub phase: Option, +} + +impl From for ChatCompletionMessage { + fn from(value: FinishReason) -> Self { + ChatCompletionMessage::default().finish_reason(value) + } +} + +/// Represents partial or full content of a message +#[derive(Clone, Debug, PartialEq, Eq, From)] +pub enum Content { + Part(ContentPart), + Full(ContentFull), +} + +impl Content { + pub fn as_str(&self) -> &str { + match self { + Content::Part(part) => &part.0, + Content::Full(full) => &full.0, + } + } + + pub fn part(content: impl ToString) -> Self { + Content::Part(ContentPart(content.to_string())) + } + + pub fn full(content: impl ToString) -> Self { + Content::Full(ContentFull(content.to_string())) + } + + pub fn is_empty(&self) -> bool { + self.as_str().is_empty() + } + + pub fn is_part(&self) -> bool { + matches!(self, Content::Part(_)) + } +} + +/// Used typically when streaming is enabled +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct ContentPart(String); + +/// Used typically when full responses are enabled (Streaming is disabled) +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct ContentFull(String); + +impl> From for Content { + fn from(value: T) -> Self { + Content::Full(ContentFull(value.as_ref().to_string())) + } +} + +/// The reason why the model stopped generating output. +/// Read more: https://platform.openai.com/docs/guides/function-calling#edge-cases +#[derive(Clone, Debug, Deserialize, Serialize, EnumString, IntoStaticStr, PartialEq, Eq)] +pub enum FinishReason { + /// The model stopped generating output because it reached the maximum + /// allowed length. + #[strum(serialize = "length")] + Length, + /// The model stopped generating output because it encountered content that + /// violated filters. + #[strum(serialize = "content_filter")] + ContentFilter, + /// The model stopped generating output because it made a tool call. + #[strum(serialize = "tool_calls")] + ToolCalls, + /// The model stopped generating output normally. + #[strum(serialize = "stop", serialize = "end_turn")] + Stop, +} + +impl ChatCompletionMessage { + pub fn assistant(content: impl Into) -> ChatCompletionMessage { + ChatCompletionMessage::default().content(content.into()) + } + + pub fn add_reasoning_detail(mut self, detail: impl Into) -> Self { + let detail = detail.into(); + if let Some(ref mut details) = self.reasoning_details { + details.push(detail); + } else { + self.reasoning_details = Some(vec![detail]); + } + self + } + + pub fn add_tool_call(mut self, call_tool: impl Into) -> Self { + self.tool_calls.push(call_tool.into()); + self + } + + pub fn extend_calls(mut self, calls: Vec>) -> Self { + self.tool_calls.extend(calls.into_iter().map(Into::into)); + self + } + + pub fn finish_reason_opt(mut self, reason: Option) -> Self { + self.finish_reason = reason; + self + } + + pub fn content_part(mut self, content: impl ToString) -> Self { + self.content = Some(Content::Part(ContentPart(content.to_string()))); + self + } + + pub fn content_full(mut self, content: impl ToString) -> Self { + self.content = Some(Content::Full(ContentFull(content.to_string()))); + self + } +} + +/// Represents a complete message from the LLM provider with all content +/// collected This is typically used after processing a stream of +/// ChatCompletionMessage +#[derive(Clone, Debug, PartialEq)] +pub struct ChatCompletionMessageFull { + pub content: String, + pub thought_signature: Option, + pub reasoning: Option, + pub tool_calls: Vec, + pub reasoning_details: Option>, + pub usage: Usage, + pub finish_reason: Option, + /// Phase label for the assistant message (e.g. `Commentary` or + /// `FinalAnswer`). + pub phase: Option, +} + +#[cfg(test)] +mod tests { + use std::str::FromStr; + + use pretty_assertions::assert_eq; + + use super::*; + #[test] + fn test_usage_accumulate_with_both_costs() { + let fixture_usage_1 = Usage { + prompt_tokens: TokenCount::Actual(100), + completion_tokens: TokenCount::Actual(50), + total_tokens: TokenCount::Actual(150), + cached_tokens: TokenCount::Actual(20), + cost: Some(0.01), + }; + + let fixture_usage_2 = Usage { + prompt_tokens: TokenCount::Actual(200), + completion_tokens: TokenCount::Actual(75), + total_tokens: TokenCount::Actual(275), + cached_tokens: TokenCount::Actual(30), + cost: Some(0.02), + }; + + let actual = fixture_usage_1.accumulate(&fixture_usage_2); + + let expected = Usage { + prompt_tokens: TokenCount::Actual(300), + completion_tokens: TokenCount::Actual(125), + total_tokens: TokenCount::Actual(425), + cached_tokens: TokenCount::Actual(50), + cost: Some(0.03), + }; + + assert_eq!(actual, expected); + } + + #[test] + fn test_usage_accumulate_mixed_token_types() { + let fixture_usage_1 = Usage { + prompt_tokens: TokenCount::Actual(100), + completion_tokens: TokenCount::Approx(50), + total_tokens: TokenCount::Actual(150), + cached_tokens: TokenCount::Actual(20), + cost: Some(0.01), + }; + + let fixture_usage_2 = Usage { + prompt_tokens: TokenCount::Approx(200), + completion_tokens: TokenCount::Actual(75), + total_tokens: TokenCount::Approx(275), + cached_tokens: TokenCount::Approx(30), + cost: Some(0.02), + }; + + let actual = fixture_usage_1.accumulate(&fixture_usage_2); + + let expected = Usage { + prompt_tokens: TokenCount::Approx(300), + completion_tokens: TokenCount::Approx(125), + total_tokens: TokenCount::Approx(425), + cached_tokens: TokenCount::Approx(50), + cost: Some(0.03), + }; + + assert_eq!(actual, expected); + } + + #[test] + fn test_usage_accumulate_partial_costs() { + let fixture_usage_1 = Usage { + prompt_tokens: TokenCount::Actual(100), + completion_tokens: TokenCount::Actual(50), + total_tokens: TokenCount::Actual(150), + cached_tokens: TokenCount::Actual(20), + cost: Some(0.01), + }; + + let fixture_usage_2 = Usage { + prompt_tokens: TokenCount::Actual(200), + completion_tokens: TokenCount::Actual(75), + total_tokens: TokenCount::Actual(275), + cached_tokens: TokenCount::Actual(30), + cost: None, + }; + + let actual = fixture_usage_1.accumulate(&fixture_usage_2); + + let expected = Usage { + prompt_tokens: TokenCount::Actual(300), + completion_tokens: TokenCount::Actual(125), + total_tokens: TokenCount::Actual(425), + cached_tokens: TokenCount::Actual(50), + cost: Some(0.01), + }; + + assert_eq!(actual, expected); + } + + #[test] + fn test_usage_accumulate_no_costs() { + let fixture_usage_1 = Usage { + prompt_tokens: TokenCount::Actual(100), + completion_tokens: TokenCount::Actual(50), + total_tokens: TokenCount::Actual(150), + cached_tokens: TokenCount::Actual(20), + cost: None, + }; + + let fixture_usage_2 = Usage { + prompt_tokens: TokenCount::Actual(200), + completion_tokens: TokenCount::Actual(75), + total_tokens: TokenCount::Actual(275), + cached_tokens: TokenCount::Actual(30), + cost: None, + }; + + let actual = fixture_usage_1.accumulate(&fixture_usage_2); + + let expected = Usage { + prompt_tokens: TokenCount::Actual(300), + completion_tokens: TokenCount::Actual(125), + total_tokens: TokenCount::Actual(425), + cached_tokens: TokenCount::Actual(50), + cost: None, + }; + + assert_eq!(actual, expected); + } + + #[test] + fn test_usage_accumulate_with_defaults() { + let fixture_usage_1 = Usage::default(); + + let fixture_usage_2 = Usage { + prompt_tokens: TokenCount::Actual(200), + completion_tokens: TokenCount::Actual(75), + total_tokens: TokenCount::Actual(275), + cached_tokens: TokenCount::Actual(30), + cost: Some(0.05), + }; + + let actual = fixture_usage_1.accumulate(&fixture_usage_2); + + let expected = Usage { + prompt_tokens: TokenCount::Actual(200), + completion_tokens: TokenCount::Actual(75), + total_tokens: TokenCount::Actual(275), + cached_tokens: TokenCount::Actual(30), + cost: Some(0.05), + }; + + assert_eq!(actual, expected); + } + + #[test] + fn test_finish_reason_from_str() { + assert_eq!( + FinishReason::from_str("length").unwrap(), + FinishReason::Length + ); + assert_eq!( + FinishReason::from_str("content_filter").unwrap(), + FinishReason::ContentFilter + ); + assert_eq!( + FinishReason::from_str("tool_calls").unwrap(), + FinishReason::ToolCalls + ); + assert_eq!(FinishReason::from_str("stop").unwrap(), FinishReason::Stop); + assert_eq!( + FinishReason::from_str("end_turn").unwrap(), + FinishReason::Stop + ); + } + + #[test] + fn test_usage_merge_anthropic_cumulative() { + // Fixture: Simulates Anthropic's message_start + message_delta pattern + // where output_tokens in message_delta is CUMULATIVE (total), not a delta. + let fixture_message_start = Usage { + prompt_tokens: TokenCount::Actual(1000), + completion_tokens: TokenCount::Actual(1), // Initial output token + total_tokens: TokenCount::Actual(1001), + cached_tokens: TokenCount::Actual(300), + cost: None, + }; + + let fixture_message_delta = Usage { + prompt_tokens: TokenCount::Actual(0), + completion_tokens: TokenCount::Actual(75), // Cumulative total, NOT delta + total_tokens: TokenCount::Actual(75), + cached_tokens: TokenCount::Actual(0), + cost: None, + }; + + let actual = fixture_message_start.merge(&fixture_message_delta); + + let expected = Usage { + prompt_tokens: TokenCount::Actual(1000), // max(1000, 0) + completion_tokens: TokenCount::Actual(75), // max(1, 75) = 75, NOT 1+75=76 + total_tokens: TokenCount::Actual(1001), // max(1001, 75) + cached_tokens: TokenCount::Actual(300), // max(300, 0) + cost: None, + }; + + assert_eq!(actual, expected); + } + + #[test] + fn test_usage_merge_preserves_costs() { + let fixture_usage_1 = Usage { + prompt_tokens: TokenCount::Actual(100), + completion_tokens: TokenCount::Actual(0), + total_tokens: TokenCount::Actual(100), + cached_tokens: TokenCount::Actual(0), + cost: Some(0.01), + }; + + let fixture_usage_2 = Usage { + prompt_tokens: TokenCount::Actual(0), + completion_tokens: TokenCount::Actual(50), + total_tokens: TokenCount::Actual(50), + cached_tokens: TokenCount::Actual(0), + cost: Some(0.02), + }; + + let actual = fixture_usage_1.merge(&fixture_usage_2); + + let expected = Usage { + prompt_tokens: TokenCount::Actual(100), + completion_tokens: TokenCount::Actual(50), + total_tokens: TokenCount::Actual(100), + cached_tokens: TokenCount::Actual(0), + cost: Some(0.03), // Costs are summed, not maxed + }; + + assert_eq!(actual, expected); + } +} diff --git a/crates/forge_domain/src/message_pattern.rs b/crates/forge_domain/src/message_pattern.rs new file mode 100644 index 0000000000000000000000000000000000000000..b9b9248a3db5f67e526cafe1cc211c33517d99bb --- /dev/null +++ b/crates/forge_domain/src/message_pattern.rs @@ -0,0 +1,262 @@ +use serde_json::json; + +use crate::{ + Context, ContextMessage, MessageEntry, ModelId, ToolCallFull, ToolCallId, ToolName, ToolResult, +}; + +/// Converts a condensed string pattern into a Context with messages. +/// +/// This utility type is primarily used in tests to quickly create Context +/// objects with specific message sequences without verbose setup code. +/// +/// # Pattern Format +/// +/// Each character in the pattern represents a message with a specific role: +/// - `'u'` = User message +/// - `'a'` = Assistant message +/// - `'s'` = System message +/// - `'t'` = Assistant message with tool call +/// - `'r'` = Tool result message +/// +/// # Examples +/// +/// ```rust,ignore +/// // Creates: User -> Assistant -> User +/// let context = MessagePattern::new("uau").build(); +/// +/// // Creates: System -> System -> User -> System -> User -> System -> User -> System -> Assistant -> Assistant -> System -> Assistant +/// let context = MessagePattern::new("ssusususaasa").build(); +/// +/// // Creates: User -> Assistant with tool call -> Tool result -> User +/// let context = MessagePattern::new("utru").build(); +/// ``` +#[derive(Debug, Clone)] +pub struct MessagePattern { + pattern: String, +} + +impl MessagePattern { + /// Creates a new MessagePattern from the given pattern string. + /// + /// # Arguments + /// + /// * `pattern` - A string where each character represents a message role: + /// - `'u'` for User + /// - `'a'` for Assistant + /// - `'s'` for System + /// - `'t'` for Assistant with tool call + /// - `'r'` for Tool result + pub fn new(pattern: impl Into) -> Self { + Self { pattern: pattern.into() } + } + + /// Builds a Context from the pattern. + /// + /// Each message will have content in the format "Message {index}" where + /// index starts from 1. Tool calls and tool results use predefined test + /// data. + /// + /// # Panics + /// + /// Panics if the pattern contains any character other than 'u', 'a', 's', + /// 't', or 'r'. + pub fn build(self) -> Context { + let model_id = ModelId::new("gpt-4"); + + let tool_call = ToolCallFull { + name: ToolName::new("read"), + call_id: Some(ToolCallId::new("call_123")), + arguments: json!({"path": "/test/path"}).into(), + thought_signature: None, + }; + + let tool_result = ToolResult::new(ToolName::new("read")) + .call_id(ToolCallId::new("call_123")) + .success(json!({"content": "File content"}).to_string()); + + let messages: Vec = self + .pattern + .chars() + .enumerate() + .map(|(i, c)| { + let content = format!("Message {}", i + 1); + match c { + 'u' => ContextMessage::user(&content, Some(model_id.clone())), + 'a' => ContextMessage::assistant(&content, None, None, None), + 's' => ContextMessage::system(&content), + 't' => ContextMessage::assistant( + &content, + None, + None, + Some(vec![tool_call.clone()]), + ), + 'r' => ContextMessage::tool_result(tool_result.clone()), + _ => { + panic!("Invalid character '{c}' in pattern. Use 'u', 'a', 's', 't', or 'r'") + } + } + }) + .map(MessageEntry::from) + .collect(); + Context::default().messages(messages) + } +} + +impl From<&str> for MessagePattern { + fn from(pattern: &str) -> Self { + Self::new(pattern) + } +} + +impl From for MessagePattern { + fn from(pattern: String) -> Self { + Self::new(pattern) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + use crate::{ContextMessage, ModelId, Role, TextMessage}; + + #[test] + fn test_message_pattern_single_user() { + let fixture = MessagePattern::new("u"); + let actual = fixture.build(); + let expected = Context::default().messages(vec![ + ContextMessage::Text( + TextMessage::new(Role::User, "Message 1").model(ModelId::new("gpt-4")), + ) + .into(), + ]); + assert_eq!(actual, expected); + } + + #[test] + fn test_message_pattern_user_assistant_user() { + let fixture = MessagePattern::new("uau"); + let actual = fixture.build(); + let expected = Context::default().messages(vec![ + ContextMessage::Text( + TextMessage::new(Role::User, "Message 1").model(ModelId::new("gpt-4")), + ) + .into(), + ContextMessage::Text(TextMessage::new(Role::Assistant, "Message 2")).into(), + ContextMessage::Text( + TextMessage::new(Role::User, "Message 3").model(ModelId::new("gpt-4")), + ) + .into(), + ]); + assert_eq!(actual, expected); + } + + #[test] + fn test_message_pattern_complex() { + let fixture = MessagePattern::new("ssusususaasa"); + let actual = fixture.build(); + + assert_eq!(actual.messages.len(), 12); + assert!(actual.messages[0].has_role(Role::System)); + assert!(actual.messages[1].has_role(Role::System)); + assert!(actual.messages[2].has_role(Role::User)); + assert!(actual.messages[3].has_role(Role::System)); + assert!(actual.messages[4].has_role(Role::User)); + assert!(actual.messages[5].has_role(Role::System)); + assert!(actual.messages[6].has_role(Role::User)); + assert!(actual.messages[7].has_role(Role::System)); + assert!(actual.messages[8].has_role(Role::Assistant)); + assert!(actual.messages[9].has_role(Role::Assistant)); + assert!(actual.messages[10].has_role(Role::System)); + assert!(actual.messages[11].has_role(Role::Assistant)); + } + + #[test] + fn test_message_pattern_empty() { + let fixture = MessagePattern::new(""); + let actual = fixture.build(); + let expected = Context::default(); + assert_eq!(actual, expected); + } + + #[test] + fn test_message_pattern_all_system() { + let fixture = MessagePattern::new("sss"); + let actual = fixture.build(); + + assert_eq!(actual.messages.len(), 3); + assert!(actual.messages.iter().all(|m| m.has_role(Role::System))); + } + + #[test] + #[should_panic(expected = "Invalid character 'x' in pattern. Use 'u', 'a', 's', 't', or 'r'")] + fn test_message_pattern_invalid_character() { + let fixture = MessagePattern::new("uax"); + fixture.build(); + } + + #[test] + fn test_message_pattern_from_str() { + let fixture = MessagePattern::from("ua"); + let actual = fixture.build(); + assert_eq!(actual.messages.len(), 2); + } + + #[test] + fn test_message_pattern_from_string() { + let fixture = MessagePattern::from("ua".to_string()); + let actual = fixture.build(); + assert_eq!(actual.messages.len(), 2); + } + + #[test] + fn test_message_pattern_content_numbering() { + let fixture = MessagePattern::new("uau"); + let actual = fixture.build(); + + assert_eq!(actual.messages[0].content().unwrap(), "Message 1"); + assert_eq!(actual.messages[1].content().unwrap(), "Message 2"); + assert_eq!(actual.messages[2].content().unwrap(), "Message 3"); + } + + #[test] + fn test_message_pattern_with_tool_call() { + let fixture = MessagePattern::new("utr"); + let actual = fixture.build(); + + assert_eq!(actual.messages.len(), 3); + assert!(actual.messages[0].has_role(Role::User)); + assert!(actual.messages[1].has_role(Role::Assistant)); + assert!(actual.messages[1].has_tool_call()); + assert!(actual.messages[2].has_tool_result()); + } + + #[test] + fn test_message_pattern_with_multiple_tool_calls() { + let fixture = MessagePattern::new("utrtr"); + let actual = fixture.build(); + + assert_eq!(actual.messages.len(), 5); + assert!(actual.messages[1].has_tool_call()); + assert!(actual.messages[2].has_tool_result()); + assert!(actual.messages[3].has_tool_call()); + assert!(actual.messages[4].has_tool_result()); + } + + #[test] + fn test_message_pattern_complex_with_tools() { + let fixture = MessagePattern::new("sutruaua"); + let actual = fixture.build(); + + assert_eq!(actual.messages.len(), 8); + assert!(actual.messages[0].has_role(Role::System)); + assert!(actual.messages[1].has_role(Role::User)); + assert!(actual.messages[2].has_tool_call()); + assert!(actual.messages[3].has_tool_result()); + assert!(actual.messages[4].has_role(Role::User)); + assert!(actual.messages[5].has_role(Role::Assistant)); + assert!(actual.messages[6].has_role(Role::User)); + assert!(actual.messages[7].has_role(Role::Assistant)); + } +} diff --git a/crates/forge_domain/src/model_config.rs b/crates/forge_domain/src/model_config.rs new file mode 100644 index 0000000000000000000000000000000000000000..930f5149690ca1a793c0172756bb05c4ec3c593b --- /dev/null +++ b/crates/forge_domain/src/model_config.rs @@ -0,0 +1,30 @@ +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +use crate::{ModelId, ProviderId}; + +/// Domain-level configuration that pairs a provider with a model. +/// +/// Used as the unified payload for [`super::ConfigOperation`] variants that +/// configure a provider/model pair (session, commit, suggest). Both fields are +/// required; use `Option` at the call-site when the configuration +/// is optional. +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)] +pub struct ModelConfig { + /// The provider ID (e.g. `"anthropic"`). + pub provider: ProviderId, + + /// The model ID to use with this provider. + pub model: ModelId, +} + +impl ModelConfig { + /// Creates a new [`ModelConfig`] with the given provider and model. + /// + /// # Arguments + /// * `provider` - The provider identifier + /// * `model` - The model identifier + pub fn new(provider: impl Into, model: impl Into) -> Self { + Self { provider: provider.into(), model: model.into() } + } +} diff --git a/crates/forge_domain/src/node.rs b/crates/forge_domain/src/node.rs new file mode 100644 index 0000000000000000000000000000000000000000..5f172185b86c456822b6bf8c705967537dbd32db --- /dev/null +++ b/crates/forge_domain/src/node.rs @@ -0,0 +1,506 @@ +use derive_more::Display; +use derive_setters::Setters; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::WorkspaceId; + +/// Progress events emitted during workspace indexing +#[derive(Debug, Clone, PartialEq)] +pub enum SyncProgress { + /// Sync operation is starting + Starting, + /// A new workspace was created on the server + WorkspaceCreated { + /// The ID of the newly created workspace + workspace_id: WorkspaceId, + }, + /// Discovering files in the directory + DiscoveringFiles { + /// ID of the current workspace + workspace_id: WorkspaceId, + /// Path being scanned + path: std::path::PathBuf, + }, + /// Files have been discovered in the directory + FilesDiscovered { + /// Total number of files found + count: usize, + }, + /// Comparing local files with server state + ComparingFiles { + /// Number of remote files in the workspace + remote_files: usize, + /// Number of local files being compared + local_files: usize, + }, + /// Diff computed showing breakdown of changes + DiffComputed { + /// Number of files added (new files) + added: usize, + /// Number of files deleted (orphaned on server) + deleted: usize, + /// Number of files modified (changed files) + modified: usize, + }, + /// Syncing files (deleting outdated + uploading new/changed) + Syncing { + /// Current progress + current: usize, + /// Total number of files to sync + total: usize, + }, + /// Sync operation completed successfully + Completed { + /// Total number of files in the workspace + total_files: usize, + /// Number of files that were uploaded (changed or new) + uploaded_files: usize, + /// Number of files that failed to sync + failed_files: usize, + }, +} + +impl SyncProgress { + /// Returns the progress weight (0-100) for this event. + pub fn weight(&self) -> Option { + match self { + Self::Syncing { current, total } => { + let sync_progress = if *total > 0 { + (*current as f64) / (*total as f64) * 100.0 + } else { + 0.0 + }; + Some(sync_progress as u64) + } + _ => None, + } + } +} + +/// Stored authentication token for the indexing service (no expiry) +/// +/// Associates a user with their indexing service authentication token +/// obtained from the remote authentication API. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkspaceAuth { + /// User ID that owns this authentication + pub user_id: UserId, + /// Authentication token (obtained from HTTP API) + pub token: crate::ApiKey, + /// When this token was stored locally + pub created_at: chrono::DateTime, +} + +impl From for crate::AuthDetails { + fn from(auth: WorkspaceAuth) -> Self { + crate::AuthDetails::ApiKey(auth.token) + } +} + +impl WorkspaceAuth { + /// Create a new indexing auth record + pub fn new(user_id: UserId, token: crate::ApiKey) -> Self { + Self { user_id, token, created_at: chrono::Utc::now() } + } +} + +/// File content for upload to workspace server +/// +/// Contains the file path (relative to workspace root) and its textual content +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FileRead { + /// File path (relative to workspace root) + pub path: String, + /// File content as UTF-8 text + pub content: String, +} + +impl FileRead { + /// Create a new file read entry + pub fn new(path: String, content: String) -> Self { + Self { path, content } + } +} + +/// Generic wrapper for workspace operations +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodeBase { + pub user_id: UserId, + pub workspace_id: WorkspaceId, + pub data: T, +} + +impl CodeBase { + pub fn new(user_id: UserId, workspace_id: WorkspaceId, data: T) -> Self { + Self { user_id, workspace_id, data } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Setters)] +#[setters(strip_option, into)] +pub struct SearchParams<'a> { + pub query: &'a str, + pub limit: Option, + pub top_k: Option, + pub use_case: String, + pub starts_with: Option, + pub ends_with: Option>, +} + +impl<'a> SearchParams<'a> { + pub fn new(query: &'a str, use_case: &str) -> Self { + Self { + query, + limit: None, + top_k: None, + use_case: use_case.to_string(), + starts_with: None, + ends_with: None, + } + } +} + +pub type CodeSearchQuery<'a> = CodeBase>; +pub type FileUpload = CodeBase>; +pub type FileDeletion = CodeBase>; +pub type WorkspaceFiles = CodeBase<()>; + +/// User identifier for codebase operations. +/// +/// Unique per machine, generated once and stored in database. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Display)] +#[display("{}", _0)] +pub struct UserId(Uuid); + +impl UserId { + /// Generate a new random user ID + pub fn generate() -> Self { + Self(Uuid::new_v4()) + } + + /// Parse a user ID from a string + /// + /// # Errors + /// Returns an error if the string is not a valid UUID + pub fn from_string(s: &str) -> anyhow::Result { + Ok(Self(Uuid::parse_str(s)?)) + } +} + +/// Node identifier for code graph nodes. +/// +/// Uniquely identifies a node in the codebase graph (file chunks, files, +/// notes, tasks, etc.). +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Display)] +#[display("{}", _0)] +pub struct NodeId(String); + +impl NodeId { + /// Create a new node ID from a string + pub fn new(id: impl Into) -> Self { + Self(id.into()) + } + + /// Get the node ID as a string slice + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl From for NodeId { + fn from(s: String) -> Self { + Self(s) + } +} + +impl From<&str> for NodeId { + fn from(s: &str) -> Self { + Self(s.to_string()) + } +} + +impl AsRef for NodeId { + fn as_ref(&self) -> &str { + &self.0 + } +} + +/// Git repository information for a workspace +/// +/// Contains commit hash and branch name for version tracking +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct GitInfo { + /// Git commit hash (e.g., "abc123...") + pub commit: String, + /// Git branch name (e.g., "main", "develop") + pub branch: String, +} + +/// Information about a workspace from the server +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct WorkspaceInfo { + /// Workspace ID + pub workspace_id: WorkspaceId, + /// Working directory path + pub working_dir: String, + /// Number of nodes created + pub node_count: Option, + /// Number of relations between nodes + pub relation_count: Option, + /// Last updated timestamp + pub last_updated: Option>, + /// Workspace created time. + pub created_at: chrono::DateTime, +} + +/// Result of a codebase sync operation +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Setters)] +pub struct FileUploadResponse { + /// Workspace ID that was synced + pub workspace_id: WorkspaceId, + /// Number of files processed + pub files_processed: usize, + /// Upload statistics + pub upload_stats: FileUploadInfo, + /// Whether a new workspace was created (vs using existing) + pub is_new_workspace: bool, +} + +impl FileUploadResponse { + /// Create new sync statistics + pub fn new( + workspace_id: WorkspaceId, + files_processed: usize, + upload_stats: FileUploadInfo, + ) -> Self { + Self { + workspace_id, + files_processed, + upload_stats, + is_new_workspace: false, + } + } +} + +/// Statistics from uploading files to the codebase server +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] +pub struct FileUploadInfo { + /// Number of code nodes created + pub nodes_created: usize, + /// Number of relations created + pub relations_created: usize, +} + +impl std::ops::Add for FileUploadInfo { + type Output = Self; + + fn add(self, other: Self) -> Self { + Self { + nodes_created: self.nodes_created + other.nodes_created, + relations_created: self.relations_created + other.relations_created, + } + } +} + +impl FileUploadInfo { + /// Create new upload statistics + pub fn new(nodes_created: usize, relations_created: usize) -> Self { + Self { nodes_created, relations_created } + } +} + +/// Results for a single codebase search query +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct CodebaseQueryResult { + /// The query string that was executed + pub query: String, + /// Relevance query used for re-ranking + pub use_case: String, + /// The search results for this query + pub results: Vec, +} + +/// Results for multiple codebase search queries +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct CodebaseSearchResults { + /// Results for each query/use_case pair + pub queries: Vec, +} + +/// A search result with its similarity score +/// +/// Wraps a code node with its semantic search scores, +/// keeping the scores separate from the node data itself. +#[derive( + Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, derive_setters::Setters, +)] +#[setters(strip_option)] +pub struct Node { + /// Node identifier + pub node_id: NodeId, + /// The node data (file, chunk, note, etc.) + #[serde(flatten)] + pub node: NodeData, + /// Relevance score (most important ranking metric) + pub relevance: Option, + /// Distance score (second ranking metric, lower is better) + pub distance: Option, +} + +/// File chunk with precise line numbers +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct FileChunk { + /// File path + pub file_path: String, + /// Code content + pub content: String, + /// Start line in the file + pub start_line: u32, + /// End line in the file + pub end_line: u32, +} + +/// Full file content +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct FileNode { + /// File path + pub file_path: String, + /// File content + pub content: String, + /// SHA-256 hash of the file content + pub hash: String, +} + +/// File reference (path only, no content) +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct FileRef { + /// File path + pub file_path: String, + /// SHA-256 hash of the file content + pub file_hash: String, +} + +/// Note content +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct Note { + /// Note content + pub content: String, +} + +/// Task description +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct Task { + /// Task description + pub task: String, +} + +/// Result of a semantic search query +/// +/// Represents different types of nodes returned from the codebase service. +/// Each variant contains only the fields relevant to that node type. +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, derive_more::From)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum NodeData { + /// File chunk with precise line numbers + #[from] + FileChunk(FileChunk), + /// Full file content + #[from] + File(FileNode), + /// File reference (path only, no content) + #[from] + FileRef(FileRef), + /// Note content + #[from] + Note(Note), + /// Task description + #[from] + Task(Task), +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_user_id_roundtrip() { + let user_id = UserId::generate(); + let s = user_id.to_string(); + let parsed = UserId::from_string(&s).unwrap(); + assert_eq!(user_id, parsed); + } + + #[test] + fn test_workspace_id_roundtrip() { + let workspace_id = WorkspaceId::generate(); + let s = workspace_id.to_string(); + let parsed = WorkspaceId::from_string(&s).unwrap(); + assert_eq!(workspace_id, parsed); + } + + #[test] + fn test_search_params_with_file_extension() { + let actual = SearchParams::new("retry mechanism", "find retry logic") + .limit(10usize) + .top_k(20u32) + .ends_with(vec![".rs".to_string()]); + + let expected = SearchParams { + query: "retry mechanism", + limit: Some(10), + top_k: Some(20), + use_case: "find retry logic".to_string(), + starts_with: None, + ends_with: Some(vec![".rs".to_string()]), + }; + + assert_eq!(actual, expected); + } + + #[test] + fn test_search_params_with_multiple_file_extensions() { + let actual = SearchParams::new("retry mechanism", "find retry logic") + .limit(10usize) + .top_k(20u32) + .ends_with(vec![ + ".rs".to_string(), + ".ts".to_string(), + ".py".to_string(), + ]); + + let expected = SearchParams { + query: "retry mechanism", + limit: Some(10), + top_k: Some(20), + use_case: "find retry logic".to_string(), + starts_with: None, + ends_with: Some(vec![ + ".rs".to_string(), + ".ts".to_string(), + ".py".to_string(), + ]), + }; + + assert_eq!(actual, expected); + } + + #[test] + fn test_search_params_without_file_extension() { + let actual = SearchParams::new("auth logic", "authentication implementation").limit(5usize); + + let expected = SearchParams { + query: "auth logic", + limit: Some(5), + top_k: None, + use_case: "authentication implementation".to_string(), + starts_with: None, + ends_with: None, + }; + + assert_eq!(actual, expected); + } +} diff --git a/crates/forge_domain/src/provider.rs b/crates/forge_domain/src/provider.rs new file mode 100644 index 0000000000000000000000000000000000000000..72c914589f963b1ed4bdfbba139b8aa836655e25 --- /dev/null +++ b/crates/forge_domain/src/provider.rs @@ -0,0 +1,1013 @@ +use std::borrow::Cow; + +use derive_more::{AsRef, Deref, From}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use strum_macros::{Display, EnumString}; +use url::Url; + +use crate::{ApiKey, AuthCredential, AuthDetails, Model, Template}; + +/// Distinguishes between different categories of providers +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Display, EnumString, Default, +)] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum ProviderType { + /// LLM providers for chat completions (default for backward compatibility) + #[default] + Llm, + /// Context engine providers for code indexing and search + ContextEngine, +} + +/// --- IMPORTANT --- +/// The order of providers is important because that would be order in which the +/// providers will be resolved +#[derive( + Debug, + Clone, + PartialEq, + Eq, + Hash, + PartialOrd, + Ord, + JsonSchema, + AsRef, + Deref, + Serialize, + Deserialize, +)] +#[schemars(with = "String")] +#[serde(from = "String")] +pub struct ProviderId(Cow<'static, str>); + +impl ProviderId { + // Built-in provider constants + pub const FORGE: ProviderId = ProviderId(Cow::Borrowed("forge")); + pub const OPENAI: ProviderId = ProviderId(Cow::Borrowed("openai")); + pub const OPEN_ROUTER: ProviderId = ProviderId(Cow::Borrowed("open_router")); + pub const REQUESTY: ProviderId = ProviderId(Cow::Borrowed("requesty")); + pub const ZAI: ProviderId = ProviderId(Cow::Borrowed("zai")); + pub const ZAI_CODING: ProviderId = ProviderId(Cow::Borrowed("zai_coding")); + pub const CEREBRAS: ProviderId = ProviderId(Cow::Borrowed("cerebras")); + pub const XAI: ProviderId = ProviderId(Cow::Borrowed("xai")); + pub const ANTHROPIC: ProviderId = ProviderId(Cow::Borrowed("anthropic")); + pub const CLAUDE_CODE: ProviderId = ProviderId(Cow::Borrowed("claude_code")); + pub const VERTEX_AI: ProviderId = ProviderId(Cow::Borrowed("vertex_ai")); + pub const VERTEX_AI_ANTHROPIC: ProviderId = ProviderId(Cow::Borrowed("vertex_ai_anthropic")); + pub const BIG_MODEL: ProviderId = ProviderId(Cow::Borrowed("big_model")); + pub const AZURE: ProviderId = ProviderId(Cow::Borrowed("azure")); + pub const GITHUB_COPILOT: ProviderId = ProviderId(Cow::Borrowed("github_copilot")); + pub const OPENAI_COMPATIBLE: ProviderId = ProviderId(Cow::Borrowed("openai_compatible")); + pub const OPENAI_RESPONSES_COMPATIBLE: ProviderId = + ProviderId(Cow::Borrowed("openai_responses_compatible")); + pub const ANTHROPIC_COMPATIBLE: ProviderId = ProviderId(Cow::Borrowed("anthropic_compatible")); + pub const FORGE_SERVICES: ProviderId = ProviderId(Cow::Borrowed("forge_services")); + pub const IO_INTELLIGENCE: ProviderId = ProviderId(Cow::Borrowed("io_intelligence")); + pub const BEDROCK: ProviderId = ProviderId(Cow::Borrowed("bedrock")); + pub const MINIMAX: ProviderId = ProviderId(Cow::Borrowed("minimax")); + pub const CODEX: ProviderId = ProviderId(Cow::Borrowed("codex")); + pub const OPENCODE_ZEN: ProviderId = ProviderId(Cow::Borrowed("opencode_zen")); + pub const OPENCODE_GO: ProviderId = ProviderId(Cow::Borrowed("opencode_go")); + pub const FIREWORKS_AI: ProviderId = ProviderId(Cow::Borrowed("fireworks-ai")); + pub const FIREWORKS_AI_FIREPASS: ProviderId = + ProviderId(Cow::Borrowed("fireworks-ai-firepass")); + pub const NOVITA: ProviderId = ProviderId(Cow::Borrowed("novita")); + pub const VIVGRID: ProviderId = ProviderId(Cow::Borrowed("vivgrid")); + pub const GOOGLE_AI_STUDIO: ProviderId = ProviderId(Cow::Borrowed("google_ai_studio")); + pub const MODAL: ProviderId = ProviderId(Cow::Borrowed("modal")); + pub const ADAL: ProviderId = ProviderId(Cow::Borrowed("adal")); + pub const XIAOMI_MIMO: ProviderId = ProviderId(Cow::Borrowed("xiaomi_mimo")); + pub const NVIDIA: ProviderId = ProviderId(Cow::Borrowed("nvidia")); + pub const AMBIENT: ProviderId = ProviderId(Cow::Borrowed("ambient")); + pub const NEURALWATT: ProviderId = ProviderId(Cow::Borrowed("neuralwatt")); + pub const ORCA_ROUTER: ProviderId = ProviderId(Cow::Borrowed("orca_router")); + pub const META: ProviderId = ProviderId(Cow::Borrowed("meta")); + pub const KIMI_CODING: ProviderId = ProviderId(Cow::Borrowed("kimi_coding")); + pub const MOONSHOT: ProviderId = ProviderId(Cow::Borrowed("moonshot")); + pub const ALIBABA_TOKEN_PLAN: ProviderId = ProviderId(Cow::Borrowed("alibaba_token_plan")); + + /// Returns all built-in provider IDs + /// + /// This includes all providers defined as constants in this implementation. + pub fn built_in_providers() -> &'static [ProviderId] { + &[ + ProviderId::FORGE, + ProviderId::OPENAI, + ProviderId::OPEN_ROUTER, + ProviderId::REQUESTY, + ProviderId::ZAI, + ProviderId::ZAI_CODING, + ProviderId::CEREBRAS, + ProviderId::XAI, + ProviderId::ANTHROPIC, + ProviderId::CLAUDE_CODE, + ProviderId::VERTEX_AI, + ProviderId::VERTEX_AI_ANTHROPIC, + ProviderId::BIG_MODEL, + ProviderId::AZURE, + ProviderId::GITHUB_COPILOT, + ProviderId::OPENAI_COMPATIBLE, + ProviderId::OPENAI_RESPONSES_COMPATIBLE, + ProviderId::ANTHROPIC_COMPATIBLE, + ProviderId::FORGE_SERVICES, + ProviderId::IO_INTELLIGENCE, + ProviderId::BEDROCK, + ProviderId::MINIMAX, + ProviderId::CODEX, + ProviderId::OPENCODE_ZEN, + ProviderId::OPENCODE_GO, + ProviderId::FIREWORKS_AI, + ProviderId::FIREWORKS_AI_FIREPASS, + ProviderId::NOVITA, + ProviderId::VIVGRID, + ProviderId::GOOGLE_AI_STUDIO, + ProviderId::MODAL, + ProviderId::ADAL, + ProviderId::XIAOMI_MIMO, + ProviderId::NVIDIA, + ProviderId::AMBIENT, + ProviderId::NEURALWATT, + ProviderId::ORCA_ROUTER, + ProviderId::META, + ProviderId::KIMI_CODING, + ProviderId::MOONSHOT, + ProviderId::ALIBABA_TOKEN_PLAN, + ] + } + + /// Returns the display name for UI (UpperCamelCase with special handling + /// for acronyms). + /// + /// This converts snake_case IDs to proper display names: + /// - "openai" -> "OpenAI" + /// - "open_router" -> "OpenRouter" + /// - "xai" -> "XAI" + fn display_name(&self) -> String { + // Special cases for known providers with acronyms + match self.0.as_ref() { + "openai" => "OpenAI".to_string(), + "xai" => "XAI".to_string(), + "zai" => "ZAI".to_string(), + "vertex_ai" => "VertexAI".to_string(), + "vertex_ai_anthropic" => "VertexAIAnthropic".to_string(), + "openai_compatible" => "OpenAICompatible".to_string(), + "openai_responses_compatible" => "OpenAIResponsesCompatible".to_string(), + "io_intelligence" => "IOIntelligence".to_string(), + "minimax" => "MiniMax".to_string(), + "codex" => "Codex".to_string(), + "opencode_zen" => "OpenCode Zen".to_string(), + "opencode_go" => "OpenCode Go".to_string(), + "fireworks-ai" => "FireworksAI".to_string(), + "fireworks-ai-firepass" => "FireworksAIFirepass".to_string(), + "novita" => "Novita".to_string(), + "vivgrid" => "Vivgrid".to_string(), + "google_ai_studio" => "GoogleAIStudio".to_string(), + "modal" => "Modal".to_string(), + "adal" => "AdaL".to_string(), + "xiaomi_mimo" => "XiaomiMimo".to_string(), + "nvidia" => "NVIDIA".to_string(), + "ambient" => "Ambient".to_string(), + "neuralwatt" => "Neuralwatt".to_string(), + "orca_router" => "OrcaRouter".to_string(), + "meta" => "Meta".to_string(), + _ => { + // For other providers, use UpperCamelCase conversion + use convert_case::{Case, Casing}; + self.0.to_case(Case::UpperCamel) + } + } + } +} + +impl std::fmt::Display for ProviderId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.display_name()) + } +} + +impl std::str::FromStr for ProviderId { + type Err = std::convert::Infallible; + + fn from_str(s: &str) -> Result { + // Check if it's a built-in provider first + let provider = match s { + "forge" => ProviderId::FORGE, + "openai" => ProviderId::OPENAI, + "open_router" => ProviderId::OPEN_ROUTER, + "requesty" => ProviderId::REQUESTY, + "zai" => ProviderId::ZAI, + "zai_coding" => ProviderId::ZAI_CODING, + "cerebras" => ProviderId::CEREBRAS, + "xai" => ProviderId::XAI, + "anthropic" => ProviderId::ANTHROPIC, + "claude_code" => ProviderId::CLAUDE_CODE, + "vertex_ai" => ProviderId::VERTEX_AI, + "big_model" => ProviderId::BIG_MODEL, + "azure" => ProviderId::AZURE, + "github_copilot" => ProviderId::GITHUB_COPILOT, + "openai_compatible" => ProviderId::OPENAI_COMPATIBLE, + "openai_responses_compatible" => ProviderId::OPENAI_RESPONSES_COMPATIBLE, + "anthropic_compatible" => ProviderId::ANTHROPIC_COMPATIBLE, + "forge_services" => ProviderId::FORGE_SERVICES, + "io_intelligence" => ProviderId::IO_INTELLIGENCE, + "minimax" => ProviderId::MINIMAX, + "codex" => ProviderId::CODEX, + "opencode_go" => ProviderId::OPENCODE_GO, + "fireworks-ai" => ProviderId::FIREWORKS_AI, + "fireworks-ai-firepass" => ProviderId::FIREWORKS_AI_FIREPASS, + "novita" => ProviderId::NOVITA, + "vertex_ai_anthropic" => ProviderId::VERTEX_AI_ANTHROPIC, + "bedrock" => ProviderId::BEDROCK, + "opencode_zen" => ProviderId::OPENCODE_ZEN, + "vivgrid" => ProviderId::VIVGRID, + "google_ai_studio" => ProviderId::GOOGLE_AI_STUDIO, + "modal" => ProviderId::MODAL, + "adal" => ProviderId::ADAL, + "xiaomi_mimo" => ProviderId::XIAOMI_MIMO, + "nvidia" => ProviderId::NVIDIA, + "ambient" => ProviderId::AMBIENT, + "neuralwatt" => ProviderId::NEURALWATT, + "orca_router" => ProviderId::ORCA_ROUTER, + "meta" => ProviderId::META, + "kimi_coding" => ProviderId::KIMI_CODING, + "moonshot" => ProviderId::MOONSHOT, + "alibaba_token_plan" => ProviderId::ALIBABA_TOKEN_PLAN, + // For custom providers, use Cow::Owned to avoid memory leaks + custom => ProviderId(Cow::Owned(custom.to_string())), + }; + Ok(provider) + } +} + +impl From for ProviderId { + fn from(s: String) -> Self { + std::str::FromStr::from_str(&s).unwrap() + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub enum ProviderResponse { + OpenAI, + OpenAIResponses, + Anthropic, + Bedrock, + Google, + OpenCode, +} + +/// Represents the source of models for a provider +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum ModelSource { + /// Can be a `Url` or a `Template` + Url(T), + Hardcoded(Vec), +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct Provider { + pub id: ProviderId, + #[serde(default)] + pub provider_type: ProviderType, + pub response: Option, + pub url: T, + pub models: Option>, + pub auth_methods: Vec, + #[serde(default)] + pub url_params: Vec, + pub credential: Option, + /// Custom HTTP headers to include in API requests for this provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub custom_headers: Option>, +} + +/// Type alias for a provider with template URLs (not yet rendered) +pub type ProviderTemplate = Provider>; + +impl Provider { + pub fn is_configured(&self) -> bool { + self.credential.is_some() + } + pub fn models(&self) -> Option<&ModelSource> { + self.models.as_ref() + } +} + +impl Provider { + pub fn url(&self) -> &Url { + &self.url + } + + pub fn api_key(&self) -> Option<&ApiKey> { + self.credential + .as_ref() + .and_then(|c| match &c.auth_details { + AuthDetails::ApiKey(key) => Some(key), + _ => None, + }) + } +} + +/// Enum for viewing providers in listings where both configured and +/// unconfigured. +#[derive(Debug, Clone, PartialEq, From)] +pub enum AnyProvider { + Url(Provider), + Template(ProviderTemplate), +} + +impl AnyProvider { + /// Returns whether this provider is configured + pub fn is_configured(&self) -> bool { + match self { + AnyProvider::Url(p) => p.is_configured(), + AnyProvider::Template(p) => p.is_configured(), + } + } + + pub fn provider_type(&self) -> &ProviderType { + match self { + AnyProvider::Url(p) => &p.provider_type, + AnyProvider::Template(t) => &t.provider_type, + } + } + + pub fn id(&self) -> ProviderId { + match self { + AnyProvider::Url(p) => p.id.clone(), + AnyProvider::Template(p) => p.id.clone(), + } + } + + /// Gets the response type + pub fn response(&self) -> Option<&ProviderResponse> { + match self { + AnyProvider::Url(p) => p.response.as_ref(), + AnyProvider::Template(p) => p.response.as_ref(), + } + } + + /// Gets the URL for this provider. + /// + /// For configured providers, returns the resolved URL. For template + /// providers with no URL parameters (i.e. a hardcoded default URL in + /// provider.json), parses and returns the template string as a URL. + /// Returns `None` for template providers that require user-supplied URL + /// parameters. + pub fn url(&self) -> Option { + match self { + AnyProvider::Url(p) => Some(p.url().clone()), + AnyProvider::Template(t) if t.url_params.is_empty() => Url::parse(&t.url.template).ok(), + AnyProvider::Template(_) => None, + } + } + pub fn url_params(&self) -> &[crate::URLParamSpec] { + match self { + AnyProvider::Url(p) => &p.url_params, + AnyProvider::Template(p) => &p.url_params, + } + } + + /// Gets the authentication methods supported by this provider + pub fn auth_methods(&self) -> &[crate::AuthMethod] { + match self { + AnyProvider::Url(p) => &p.auth_methods, + AnyProvider::Template(p) => &p.auth_methods, + } + } + + /// Consumes self and returns the configured provider if this is a URL + /// provider with credentials + pub fn into_configured(self) -> Option> { + match self { + AnyProvider::Url(p) if p.is_configured() => Some(p), + _ => None, + } + } +} + +/// Represents a provider with its available models +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ProviderModels { + /// The provider identifier + pub provider_id: ProviderId, + /// Available models from this provider + pub models: Vec, +} + +#[cfg(test)] +mod test_helpers { + use std::collections::HashMap; + + use super::*; + + fn make_credential(provider_id: ProviderId, key: &str) -> Option { + Some(AuthCredential { + id: provider_id, + auth_details: AuthDetails::ApiKey(ApiKey::from(key.to_string())), + url_params: HashMap::new(), + }) + } + + /// Test helper for creating a ZAI provider + pub(super) fn zai(key: &str) -> Provider { + Provider { + id: ProviderId::ZAI, + provider_type: Default::default(), + response: Some(ProviderResponse::OpenAI), + url: Url::parse("https://api.z.ai/api/paas/v4/chat/completions").unwrap(), + auth_methods: vec![crate::AuthMethod::ApiKey], + url_params: vec![], + credential: make_credential(ProviderId::ZAI, key), + custom_headers: None, + models: Some(ModelSource::Url( + Url::parse("https://api.z.ai/api/paas/v4/models").unwrap(), + )), + } + } + + /// Test helper for creating a ZAI Coding provider + pub(super) fn zai_coding(key: &str) -> Provider { + Provider { + id: ProviderId::ZAI_CODING, + provider_type: Default::default(), + response: Some(ProviderResponse::OpenAI), + url: Url::parse("https://api.z.ai/api/coding/paas/v4/chat/completions").unwrap(), + auth_methods: vec![crate::AuthMethod::ApiKey], + url_params: vec![], + credential: make_credential(ProviderId::ZAI_CODING, key), + custom_headers: None, + models: Some(ModelSource::Url( + Url::parse("https://api.z.ai/api/paas/v4/models").unwrap(), + )), + } + } + + /// Test helper for creating an OpenAI provider + pub(super) fn openai(key: &str) -> Provider { + Provider { + id: ProviderId::OPENAI, + provider_type: Default::default(), + response: Some(ProviderResponse::OpenAI), + url: Url::parse("https://api.openai.com/v1/chat/completions").unwrap(), + auth_methods: vec![crate::AuthMethod::ApiKey], + url_params: vec![], + credential: make_credential(ProviderId::OPENAI, key), + custom_headers: None, + models: Some(ModelSource::Url( + Url::parse("https://api.openai.com/v1/models").unwrap(), + )), + } + } + + /// Test helper for creating an XAI provider + pub(super) fn xai(key: &str) -> Provider { + Provider { + id: ProviderId::XAI, + provider_type: Default::default(), + response: Some(ProviderResponse::OpenAI), + url: Url::parse("https://api.x.ai/v1/chat/completions").unwrap(), + auth_methods: vec![crate::AuthMethod::ApiKey], + url_params: vec![], + credential: make_credential(ProviderId::XAI, key), + custom_headers: None, + models: Some(ModelSource::Url( + Url::parse("https://api.x.ai/v1/models").unwrap(), + )), + } + } + + /// Test helper for creating a Vertex AI provider + pub(super) fn vertex_ai(key: &str, project_id: &str, location: &str) -> Provider { + let (chat_url, model_url) = if location == "global" { + ( + format!( + "https://aiplatform.googleapis.com/v1/projects/{}/locations/{}/endpoints/openapi/chat/completions", + project_id, location + ), + format!( + "https://aiplatform.googleapis.com/v1/projects/{}/locations/{}/endpoints/openapi/models", + project_id, location + ), + ) + } else { + ( + format!( + "https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/endpoints/openapi/chat/completions", + location, project_id, location + ), + format!( + "https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/endpoints/openapi/models", + location, project_id, location + ), + ) + }; + Provider { + id: ProviderId::VERTEX_AI, + provider_type: Default::default(), + response: Some(ProviderResponse::OpenAI), + url: Url::parse(&chat_url).unwrap(), + auth_methods: vec![crate::AuthMethod::ApiKey], + url_params: ["project_id", "location"] + .iter() + .map(|&s| s.to_string().into()) + .collect(), + credential: make_credential(ProviderId::VERTEX_AI, key), + custom_headers: None, + models: Some(ModelSource::Url(Url::parse(&model_url).unwrap())), + } + } + + /// Test helper for creating an IO Intelligence provider + pub(super) fn io_intelligence(key: &str) -> Provider { + Provider { + id: ProviderId::IO_INTELLIGENCE, + provider_type: Default::default(), + response: Some(ProviderResponse::OpenAI), + url: Url::parse("https://api.intelligence.io.solutions/api/v1/chat/completions") + .unwrap(), + auth_methods: vec![crate::AuthMethod::ApiKey], + url_params: vec![], + credential: make_credential(ProviderId::IO_INTELLIGENCE, key), + custom_headers: None, + models: Some(ModelSource::Url( + Url::parse("https://api.intelligence.io.solutions/api/v1/models").unwrap(), + )), + } + } + + /// Test helper for creating an Azure provider + pub(super) fn azure( + key: &str, + resource_name: &str, + deployment_name: &str, + api_version: &str, + ) -> Provider { + let chat_url = format!( + "https://{}.openai.azure.com/openai/deployments/{}/chat/completions?api-version={}", + resource_name, deployment_name, api_version + ); + let model_url = format!( + "https://{}.openai.azure.com/openai/models?api-version={}", + resource_name, api_version + ); + + Provider { + id: ProviderId::AZURE, + provider_type: Default::default(), + response: Some(ProviderResponse::OpenAI), + url: Url::parse(&chat_url).unwrap(), + auth_methods: vec![crate::AuthMethod::ApiKey], + url_params: ["resource_name", "deployment_name", "api_version"] + .iter() + .map(|&s| s.to_string().into()) + .collect(), + credential: make_credential(ProviderId::AZURE, key), + custom_headers: None, + models: Some(ModelSource::Url(Url::parse(&model_url).unwrap())), + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::str::FromStr; + + use pretty_assertions::assert_eq; + + use super::test_helpers::*; + use super::*; + + #[test] + fn test_provider_id_display_name() { + assert_eq!(ProviderId::OPENAI.to_string(), "OpenAI"); + assert_eq!(ProviderId::OPEN_ROUTER.to_string(), "OpenRouter"); + assert_eq!(ProviderId::ZAI.to_string(), "ZAI"); + assert_eq!(ProviderId::XAI.to_string(), "XAI"); + assert_eq!(ProviderId::ANTHROPIC.to_string(), "Anthropic"); + assert_eq!(ProviderId::GITHUB_COPILOT.to_string(), "GithubCopilot"); + assert_eq!(ProviderId::VERTEX_AI.to_string(), "VertexAI"); + assert_eq!( + ProviderId::OPENAI_COMPATIBLE.to_string(), + "OpenAICompatible" + ); + assert_eq!( + ProviderId::OPENAI_RESPONSES_COMPATIBLE.to_string(), + "OpenAIResponsesCompatible" + ); + assert_eq!( + ProviderId::ANTHROPIC_COMPATIBLE.to_string(), + "AnthropicCompatible" + ); + assert_eq!(ProviderId::IO_INTELLIGENCE.to_string(), "IOIntelligence"); + assert_eq!(ProviderId::CODEX.to_string(), "Codex"); + assert_eq!(ProviderId::FIREWORKS_AI.to_string(), "FireworksAI"); + assert_eq!(ProviderId::VIVGRID.to_string(), "Vivgrid"); + assert_eq!(ProviderId::OPENCODE_ZEN.to_string(), "OpenCode Zen"); + assert_eq!(ProviderId::OPENCODE_GO.to_string(), "OpenCode Go"); + assert_eq!(ProviderId::GOOGLE_AI_STUDIO.to_string(), "GoogleAIStudio"); + assert_eq!(ProviderId::NVIDIA.to_string(), "NVIDIA"); + assert_eq!(ProviderId::AMBIENT.to_string(), "Ambient"); + assert_eq!(ProviderId::ORCA_ROUTER.to_string(), "OrcaRouter"); + assert_eq!(ProviderId::META.to_string(), "Meta"); + } + + #[test] + fn test_codex_from_str() { + let actual = ProviderId::from_str("codex").unwrap(); + let expected = ProviderId::CODEX; + assert_eq!(actual, expected); + } + + #[test] + fn test_fireworks_ai_from_str() { + let actual = ProviderId::from_str("fireworks-ai").unwrap(); + let expected = ProviderId::FIREWORKS_AI; + assert_eq!(actual, expected); + } + + #[test] + fn test_vivgrid_from_str() { + let actual = ProviderId::from_str("vivgrid").unwrap(); + let expected = ProviderId::VIVGRID; + assert_eq!(actual, expected); + } + + #[test] + fn test_opencode_go_from_str() { + let actual = ProviderId::from_str("opencode_go").unwrap(); + let expected = ProviderId::OPENCODE_GO; + assert_eq!(actual, expected); + } + + #[test] + fn test_codex_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::CODEX)); + assert!(built_in.contains(&ProviderId::OPENAI_RESPONSES_COMPATIBLE)); + assert!(built_in.contains(&ProviderId::FIREWORKS_AI)); + assert!(built_in.contains(&ProviderId::VIVGRID)); + assert!(built_in.contains(&ProviderId::OPENCODE_GO)); + assert!(built_in.contains(&ProviderId::GOOGLE_AI_STUDIO)); + assert!(built_in.contains(&ProviderId::NVIDIA)); + assert!(built_in.contains(&ProviderId::AMBIENT)); + assert!(built_in.contains(&ProviderId::ORCA_ROUTER)); + assert!(built_in.contains(&ProviderId::META)); + } + + #[test] + fn test_google_ai_studio_from_str() { + let actual = ProviderId::from_str("google_ai_studio").unwrap(); + let expected = ProviderId::GOOGLE_AI_STUDIO; + assert_eq!(actual, expected); + } + + #[test] + fn test_modal_from_str() { + let actual = ProviderId::from_str("modal").unwrap(); + let expected = ProviderId::MODAL; + assert_eq!(actual, expected); + } + + #[test] + fn test_adal_from_str() { + let actual = ProviderId::from_str("adal").unwrap(); + let expected = ProviderId::ADAL; + assert_eq!(actual, expected); + } + + #[test] + fn test_modal_display_name() { + assert_eq!(ProviderId::MODAL.to_string(), "Modal"); + } + + #[test] + fn test_modal_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::MODAL)); + } + + #[test] + fn test_adal_display_name() { + assert_eq!(ProviderId::ADAL.to_string(), "AdaL"); + } + + #[test] + fn test_adal_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::ADAL)); + } + + #[test] + fn test_xiaomi_mimo_from_str() { + let actual = ProviderId::from_str("xiaomi_mimo").unwrap(); + let expected = ProviderId::XIAOMI_MIMO; + assert_eq!(actual, expected); + } + + #[test] + fn test_xiaomi_mimo_display_name() { + assert_eq!(ProviderId::XIAOMI_MIMO.to_string(), "XiaomiMimo"); + } + + #[test] + fn test_xiaomi_mimo_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::XIAOMI_MIMO)); + } + + #[test] + fn test_ambient_from_str() { + let actual = ProviderId::from_str("ambient").unwrap(); + let expected = ProviderId::AMBIENT; + assert_eq!(actual, expected); + } + + #[test] + fn test_ambient_display_name() { + assert_eq!(ProviderId::AMBIENT.to_string(), "Ambient"); + } + + #[test] + fn test_ambient_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::AMBIENT)); + } + + #[test] + fn test_neuralwatt_from_str() { + let actual = ProviderId::from_str("neuralwatt").unwrap(); + let expected = ProviderId::NEURALWATT; + assert_eq!(actual, expected); + } + + #[test] + fn test_neuralwatt_display_name() { + assert_eq!(ProviderId::NEURALWATT.to_string(), "Neuralwatt"); + } + + #[test] + fn test_neuralwatt_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::NEURALWATT)); + } + + #[test] + fn test_orca_router_from_str() { + let actual = ProviderId::from_str("orca_router").unwrap(); + let expected = ProviderId::ORCA_ROUTER; + assert_eq!(actual, expected); + } + + #[test] + fn test_orca_router_display_name() { + assert_eq!(ProviderId::ORCA_ROUTER.to_string(), "OrcaRouter"); + } + + #[test] + fn test_orca_router_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::ORCA_ROUTER)); + } + + #[test] + fn test_meta_from_str() { + let actual = ProviderId::from_str("meta").unwrap(); + let expected = ProviderId::META; + assert_eq!(actual, expected); + } + + #[test] + fn test_meta_display_name() { + assert_eq!(ProviderId::META.to_string(), "Meta"); + } + + #[test] + fn test_meta_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::META)); + } + + #[test] + fn test_moonshot_display_name() { + let actual = ProviderId::MOONSHOT.to_string(); + let expected = "Moonshot".to_string(); + assert_eq!(actual, expected); + } + + #[test] + fn test_moonshot_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::MOONSHOT)); + } + + #[test] + fn test_moonshot_from_str_roundtrip() { + let actual = ProviderId::from_str("moonshot").unwrap(); + let expected = ProviderId::MOONSHOT; + assert_eq!(actual, expected); + } + + #[test] + fn test_kimi_coding_from_str_roundtrip() { + let actual = ProviderId::from_str("kimi_coding").unwrap(); + let expected = ProviderId::KIMI_CODING; + assert_eq!(actual, expected); + } + + #[test] + fn test_alibaba_token_plan_display_name() { + let actual = ProviderId::ALIBABA_TOKEN_PLAN.to_string(); + let expected = "AlibabaTokenPlan".to_string(); + assert_eq!(actual, expected); + } + + #[test] + fn test_alibaba_token_plan_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::ALIBABA_TOKEN_PLAN)); + } + + #[test] + fn test_alibaba_token_plan_from_str_roundtrip() { + let actual = ProviderId::from_str("alibaba_token_plan").unwrap(); + let expected = ProviderId::ALIBABA_TOKEN_PLAN; + assert_eq!(actual, expected); + } + + #[test] + fn test_io_intelligence() { + let fixture = "test_key"; + let actual = io_intelligence(fixture); + let expected = Provider { + id: ProviderId::IO_INTELLIGENCE, + provider_type: Default::default(), + response: Some(ProviderResponse::OpenAI), + url: Url::from_str("https://api.intelligence.io.solutions/api/v1/chat/completions") + .unwrap(), + credential: Some(AuthCredential { + id: ProviderId::IO_INTELLIGENCE, + auth_details: AuthDetails::ApiKey(ApiKey::from(fixture.to_string())), + url_params: HashMap::new(), + }), + auth_methods: vec![crate::AuthMethod::ApiKey], + url_params: vec![], + models: Some(ModelSource::Url( + Url::from_str("https://api.intelligence.io.solutions/api/v1/models").unwrap(), + )), + custom_headers: None, + }; + assert_eq!(actual, expected); + } + + #[test] + fn test_xai() { + let fixture = "test_key"; + let actual = xai(fixture); + let expected = Provider { + id: ProviderId::XAI, + provider_type: Default::default(), + response: Some(ProviderResponse::OpenAI), + url: Url::from_str("https://api.x.ai/v1/chat/completions").unwrap(), + credential: Some(AuthCredential { + id: ProviderId::XAI, + auth_details: AuthDetails::ApiKey(ApiKey::from(fixture.to_string())), + url_params: HashMap::new(), + }), + auth_methods: vec![crate::AuthMethod::ApiKey], + url_params: vec![], + models: Some(ModelSource::Url( + Url::from_str("https://api.x.ai/v1/models").unwrap(), + )), + custom_headers: None, + }; + assert_eq!(actual, expected); + } + + #[test] + fn test_is_xai_with_direct_comparison() { + let fixture_xai = xai("key"); + assert_eq!(fixture_xai.id, ProviderId::XAI); + + let fixture_other = openai("key"); + assert_ne!(fixture_other.id, ProviderId::XAI); + } + + #[test] + fn test_zai_coding_to_chat_url() { + let fixture = zai_coding("test_key"); + let actual = fixture.url.clone(); + let expected = Url::parse("https://api.z.ai/api/coding/paas/v4/chat/completions").unwrap(); + assert_eq!(actual, expected); + } + + #[test] + fn test_zai_coding_to_model_url() { + let fixture = zai_coding("test_key"); + let actual = fixture.models.clone(); + let expected = Some(ModelSource::Url( + Url::parse("https://api.z.ai/api/paas/v4/models").unwrap(), + )); + assert_eq!(actual, expected); + } + + #[test] + fn test_regular_zai_to_chat_url() { + let fixture = zai("test_key"); + let actual = fixture.url.clone(); + let expected = Url::parse("https://api.z.ai/api/paas/v4/chat/completions").unwrap(); + assert_eq!(actual, expected); + } + + #[test] + fn test_regular_zai_to_model_url() { + let fixture = zai("test_key"); + let actual = fixture.models.clone(); + let expected = Some(ModelSource::Url( + Url::parse("https://api.z.ai/api/paas/v4/models").unwrap(), + )); + assert_eq!(actual, expected); + } + + #[test] + fn test_vertex_ai_global_location() { + let fixture = vertex_ai("test_token", "forge-452914", "global"); + let actual = fixture.url.clone(); + let expected = Url::parse("https://aiplatform.googleapis.com/v1/projects/forge-452914/locations/global/endpoints/openapi/chat/completions").unwrap(); + assert_eq!(actual, expected); + } + + #[test] + fn test_vertex_ai_regular_location() { + let fixture = vertex_ai("test_token", "test_project", "us-central1"); + let actual = fixture.url.clone(); + let expected = Url::parse("https://us-central1-aiplatform.googleapis.com/v1/projects/test_project/locations/us-central1/endpoints/openapi/chat/completions").unwrap(); + assert_eq!(actual, expected); + } + + #[test] + fn test_fireworks_ai_firepass_from_str() { + let actual = ProviderId::from_str("fireworks-ai-firepass").unwrap(); + let expected = ProviderId::FIREWORKS_AI_FIREPASS; + assert_eq!(actual, expected); + } + + #[test] + fn test_fireworks_ai_firepass_display_name() { + assert_eq!( + ProviderId::FIREWORKS_AI_FIREPASS.to_string(), + "FireworksAIFirepass" + ); + } + + #[test] + fn test_fireworks_ai_firepass_in_built_in_providers() { + let built_in = ProviderId::built_in_providers(); + assert!(built_in.contains(&ProviderId::FIREWORKS_AI_FIREPASS)); + } + + #[test] + fn test_azure_provider() { + let fixture = azure("test_key", "my-resource", "gpt-4", "2024-02-15-preview"); + + // Check chat completion URL (url field now contains the chat completion URL) + let actual_chat = fixture.url.clone(); + let expected_chat = Url::parse("https://my-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions?api-version=2024-02-15-preview").unwrap(); + assert_eq!(actual_chat, expected_chat); + + // Check model URL + let actual_model = fixture.models.clone(); + let expected_model = Some(ModelSource::Url( + Url::parse( + "https://my-resource.openai.azure.com/openai/models?api-version=2024-02-15-preview", + ) + .unwrap(), + )); + assert_eq!(actual_model, expected_model); + + assert_eq!(fixture.id, ProviderId::AZURE); + assert_eq!(fixture.response, Some(ProviderResponse::OpenAI)); + } + + #[test] + fn test_azure_provider_with_different_params() { + let fixture = azure("another_key", "east-us", "gpt-35-turbo", "2023-05-15"); + + // Check chat completion URL + let actual_chat = fixture.url.clone(); + let expected_chat = Url::parse("https://east-us.openai.azure.com/openai/deployments/gpt-35-turbo/chat/completions?api-version=2023-05-15").unwrap(); + assert_eq!(actual_chat, expected_chat); + + // Check model URL + let actual_model = fixture.models.clone(); + let expected_model = Some(ModelSource::Url( + Url::parse("https://east-us.openai.azure.com/openai/models?api-version=2023-05-15") + .unwrap(), + )); + assert_eq!(actual_model, expected_model); + } +} diff --git a/crates/forge_domain/src/shell.rs b/crates/forge_domain/src/shell.rs new file mode 100644 index 0000000000000000000000000000000000000000..84df33a2ccf7ec7845673567c9484e98f40e78ff --- /dev/null +++ b/crates/forge_domain/src/shell.rs @@ -0,0 +1,14 @@ +/// Output from a command execution +#[derive(Debug, Clone)] +pub struct CommandOutput { + pub command: String, + pub stdout: String, + pub stderr: String, + pub exit_code: Option, +} + +impl CommandOutput { + pub fn success(&self) -> bool { + self.exit_code.is_none_or(|code| code >= 0) + } +} diff --git a/crates/forge_domain/src/skill.rs b/crates/forge_domain/src/skill.rs new file mode 100644 index 0000000000000000000000000000000000000000..d86c27f2d69fd3dbc48df6eb482d17c889c9038f --- /dev/null +++ b/crates/forge_domain/src/skill.rs @@ -0,0 +1,102 @@ +use std::path::PathBuf; + +use derive_setters::Setters; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// Represents a reusable skill with a name, file path, and prompt content +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Setters, JsonSchema)] +#[setters(strip_option, into)] +pub struct Skill { + /// Name of the skill + pub name: String, + + /// File path to the skill markdown file + pub path: Option, + + /// Content/prompt loaded from the markdown file + pub command: String, + + /// Description of the skill + pub description: String, + + /// List of resource files in the skill directory + pub resources: Vec, +} + +impl Skill { + /// Creates a new Skill with required fields + /// + /// # Arguments + /// + /// * `name` - The name identifier for the skill + /// * `prompt` - The skill prompt content + /// * `description` - A brief description of the skill + pub fn new( + name: impl Into, + prompt: impl Into, + description: impl Into, + ) -> Self { + Self { + name: name.into(), + path: None, + command: prompt.into(), + description: description.into(), + resources: Vec::new(), + } + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_skill_creation() { + // Fixture + let fixture = Skill::new( + "code_review", + "Review this code", + "A skill for reviewing code quality", + ) + .path("/skills/code_review.md"); + + // Act + let actual = ( + fixture.name.clone(), + fixture.path.clone(), + fixture.command.clone(), + fixture.description.clone(), + ); + + // Assert + let expected = ( + "code_review".to_string(), + Some("/skills/code_review.md".into()), + "Review this code".to_string(), + "A skill for reviewing code quality".to_string(), + ); + assert_eq!(actual, expected); + } + + #[test] + fn test_skill_with_setters() { + // Fixture + let fixture = Skill::new("test", "prompt", "desc") + .path("/path") + .name("updated_name") + .path("/updated/path") + .command("updated prompt") + .description("updated description"); + + // Act + let actual = fixture; + + // Assert + let expected = Skill::new("updated_name", "updated prompt", "updated description") + .path("/updated/path"); + assert_eq!(actual, expected); + } +} diff --git a/crates/forge_domain/src/snapshot.rs b/crates/forge_domain/src/snapshot.rs new file mode 100644 index 0000000000000000000000000000000000000000..d75826526117bdbcb089ef84ea8b06749930c961 --- /dev/null +++ b/crates/forge_domain/src/snapshot.rs @@ -0,0 +1,148 @@ +use std::fmt::{Display, Formatter}; +use std::hash::Hasher; +use std::path::PathBuf; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +/// A newtype for snapshot IDs, internally using UUID +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)] +pub struct SnapshotId(Uuid); + +impl SnapshotId { + /// Create a new random SnapshotId + pub fn new() -> Self { + Self(Uuid::new_v4()) + } + + /// Parse a SnapshotId from a string + pub fn parse(s: &str) -> Option { + Uuid::parse_str(s).ok().map(Self) + } + + /// Get the underlying UUID + pub fn uuid(&self) -> &Uuid { + &self.0 + } +} + +impl Default for SnapshotId { + fn default() -> Self { + Self::new() + } +} + +impl Display for SnapshotId { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +impl From for SnapshotId { + fn from(uuid: Uuid) -> Self { + Self(uuid) + } +} + +/// Represents information about a file snapshot +/// +/// Contains details about when the snapshot was created, +/// the original file path, the snapshot location, and file size. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Snapshot { + /// Unique ID for the file + pub id: SnapshotId, + + /// Unix timestamp when the snapshot was created + pub timestamp: Duration, + + /// Original file path that is being processed + pub path: String, +} + +impl Snapshot { + pub fn create(path: PathBuf) -> anyhow::Result { + let path = match path.canonicalize() { + Ok(p) => p, + Err(_) => { + if path.is_absolute() { + path + } else { + anyhow::bail!( + "Path must be absolute. Please provide an absolute path starting with '/' (Unix) or 'C:\\' (Windows)" + ); + } + } + }; + let timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?; + + Ok(Self { + id: SnapshotId::new(), + timestamp, + path: path.display().to_string(), + }) + } + + /// Create a hash of a file path for storage + pub fn path_hash(&self) -> String { + let mut hasher = fnv_rs::Fnv64::default(); + hasher.write(self.path.as_bytes()); + format!("{:x}", hasher.finish()) + } + + /// Create a snapshot filename from a path and timestamp + pub fn snapshot_path(&self, cwd: Option) -> PathBuf { + // Convert Duration to SystemTime then to a formatted string + let datetime = UNIX_EPOCH + self.timestamp; + // Format: YYYY-MM-DD_HH-MM-SS-nnnnnnnnn (including nanoseconds) + let formatted_time = chrono::DateTime::::from(datetime) + .format("%Y-%m-%d_%H-%M-%S-%9f") + .to_string(); + + let filename = format!("{formatted_time}.snap"); + let path = PathBuf::from(self.path_hash()).join(PathBuf::from(filename)); + if let Some(cwd) = cwd { + cwd.join(path) + } else { + path + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_create_with_nonexistent_absolute_path() { + // Test with a non-existent absolute path + let nonexistent_path = PathBuf::from("/this/path/does/not/exist/file.txt"); + let snapshot = Snapshot::create(nonexistent_path.clone()).unwrap(); + + assert!(!snapshot.id.to_string().is_empty()); + assert!(snapshot.timestamp.as_secs() > 0); + // Should use the original absolute path since canonicalize fails + assert_eq!(snapshot.path, nonexistent_path.display().to_string()); + } + + #[test] + fn test_create_with_nonexistent_relative_path() { + // Test with a non-existent relative path + let nonexistent_path = PathBuf::from("nonexistent/file.txt"); + let snapshot = Snapshot::create(nonexistent_path.clone()); + assert!(snapshot.is_err()); + } + + #[cfg(windows)] + #[test] + fn test_create_with_nonexistent_absolute_windows_path() { + // Test with Windows-style absolute path that doesn't exist + let nonexistent_path = PathBuf::from("C:\\nonexistent\\windows\\path\\file.txt"); + let snapshot = Snapshot::create(nonexistent_path.clone()).unwrap(); + + assert!(!snapshot.id.to_string().is_empty()); + assert!(snapshot.timestamp.as_secs() > 0); + assert_eq!(snapshot.path, nonexistent_path.display().to_string()); + } +} diff --git a/crates/forge_domain/src/suggestion.rs b/crates/forge_domain/src/suggestion.rs new file mode 100644 index 0000000000000000000000000000000000000000..85c1e146afe868bd5ce99ec2b91b6c19be14a42e --- /dev/null +++ b/crates/forge_domain/src/suggestion.rs @@ -0,0 +1,7 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize)] +pub struct Suggestion { + pub use_case: String, + pub suggestion: String, +} diff --git a/crates/forge_domain/src/temperature.rs b/crates/forge_domain/src/temperature.rs new file mode 100644 index 0000000000000000000000000000000000000000..c7b2131435b9903889cfcd802edd3bae4153f165 --- /dev/null +++ b/crates/forge_domain/src/temperature.rs @@ -0,0 +1,215 @@ +use std::fmt; +use std::ops::Deref; + +use schemars::JsonSchema; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +/// A newtype for temperature values with built-in validation +/// +/// Temperature controls the randomness in the model's output: +/// - Lower values (e.g., 0.1) make responses more focused, deterministic, and +/// coherent +/// - Higher values (e.g., 0.8) make responses more creative, diverse, and +/// exploratory +/// - Valid range is 0.0 to 2.0 +#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, JsonSchema)] +pub struct Temperature(f32); + +impl Temperature { + /// Creates a new Temperature value, returning an error if outside the valid + /// range (0.0 to 2.0) + pub fn new(value: f32) -> Result { + if Self::is_valid(value) { + Ok(Self(value)) + } else { + Err(format!( + "temperature must be between 0.0 and 2.0, got {value}" + )) + } + } + + /// Creates a new Temperature value without validation + /// + /// # Safety + /// This function should only be used when the value is known to be valid + pub fn new_unchecked(value: f32) -> Self { + debug_assert!(Self::is_valid(value), "invalid temperature: {value}"); + Self(value) + } + + /// Returns true if the temperature value is within the valid range (0.0 to + /// 2.0) + pub fn is_valid(value: f32) -> bool { + (0.0..=2.0).contains(&value) + } + + /// Returns the inner f32 value + pub fn value(&self) -> f32 { + self.0 + } +} + +impl Deref for Temperature { + type Target = f32; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl From for f32 { + fn from(temp: Temperature) -> Self { + temp.0 + } +} + +impl From for Temperature { + fn from(value: f32) -> Self { + Temperature::new_unchecked(value) + } +} + +impl fmt::Display for Temperature { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +impl Serialize for Temperature { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + // Convert to string with fixed precision to avoid floating point issues + // and then parse back to ensure consistent serialization + let formatted = format!("{:.1}", self.0); + let value = formatted.parse::().unwrap(); + serializer.serialize_f32(value) + } +} + +impl<'de> Deserialize<'de> for Temperature { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + use serde::de::Error; + let value = f32::deserialize(deserializer)?; + if Self::is_valid(value) { + Ok(Self(value)) + } else { + Err(Error::custom(format!( + "temperature must be between 0.0 and 2.0, got {value}" + ))) + } + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use serde_json::json; + + use super::*; + + #[test] + fn test_temperature_creation() { + // Valid temperatures should be created successfully + let valid_temps = [0.0, 0.5, 1.0, 1.5, 2.0]; + for temp in valid_temps { + let result = Temperature::new(temp); + assert!(result.is_ok(), "Temperature {temp} should be valid"); + assert_eq!(result.unwrap().value(), temp); + } + + // Invalid temperatures should return an error + let invalid_temps = [-0.1, 2.1, 3.0, -1.0, 10.0]; + for temp in invalid_temps { + let result = Temperature::new(temp); + assert!(result.is_err(), "Temperature {temp} should be invalid"); + assert!( + result + .unwrap_err() + .contains("temperature must be between 0.0 and 2.0"), + "Error should mention valid range" + ); + } + } + + #[test] + fn test_temperature_serialization() { + let temp = Temperature::new(0.7).unwrap(); + let json = serde_json::to_value(temp).unwrap(); + + // When serializing floating point numbers, precision issues might occur + // So we'll check if the serialized value is approximately equal to 0.7 + if let serde_json::Value::Number(num) = &json { + let float_val = num.as_f64().unwrap(); + assert!( + (float_val - 0.7).abs() < 0.001, + "Expected approximately 0.7, got {float_val}" + ); + } else { + panic!("Expected a number, got {json:?}"); + } + } + + #[test] + fn test_temperature_deserialization() { + // Valid temperature values should deserialize correctly + let valid_temps = [0.0, 0.5, 1.0, 1.5, 2.0]; + for temp_value in valid_temps { + let json = json!(temp_value); + let temp: Result = serde_json::from_value(json); + assert!( + temp.is_ok(), + "Valid temperature {temp_value} should deserialize" + ); + assert_eq!(temp.unwrap().value(), temp_value); + } + + // Invalid temperature values should fail deserialization + let invalid_temps = [-0.1, 2.1, 3.0, -1.0, 10.0]; + for temp_value in invalid_temps { + let json = json!(temp_value); + let temp: Result = serde_json::from_value(json); + assert!( + temp.is_err(), + "Invalid temperature {temp_value} should fail deserialization" + ); + let err = temp.unwrap_err().to_string(); + assert!( + err.contains("temperature must be between 0.0 and 2.0"), + "Error should mention valid range: {err}" + ); + } + } + + #[test] + fn test_temperature_in_struct() { + #[derive(Serialize, Deserialize, Debug)] + struct TestStruct { + temp: Temperature, + } + + // Valid temperature + let json = json!({ + "temp": 0.7 + }); + let test_struct: Result = serde_json::from_value(json); + assert!(test_struct.is_ok()); + assert_eq!(test_struct.unwrap().temp.value(), 0.7); + + // Invalid temperature + let json = json!({ + "temp": 2.5 + }); + let test_struct: Result = serde_json::from_value(json); + assert!(test_struct.is_err()); + let err = test_struct.unwrap_err().to_string(); + assert!( + err.contains("temperature must be between 0.0 and 2.0"), + "Error should mention valid range: {err}" + ); + } +} diff --git a/crates/forge_domain/src/tool_order.rs b/crates/forge_domain/src/tool_order.rs new file mode 100644 index 0000000000000000000000000000000000000000..f07543a62c1a04163a9d13c5b7e8fe4ebcb542a5 --- /dev/null +++ b/crates/forge_domain/src/tool_order.rs @@ -0,0 +1,436 @@ +use std::cmp::Ordering; +use std::collections::HashMap; +use std::sync::Arc; + +use glob::Pattern; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +use crate::{ToolDefinition, ToolName}; + +/// Defines the ordering for tools in an agent's context. +/// Tools are ordered based on weights - higher weight tools appear first. +/// When the list is empty, tools are sorted alphabetically. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +pub struct ToolOrder { + /// The ordered list of tool names and patterns + tools: Vec, + /// Weight map: tool name -> weight (position index in the tools vector) + /// Pre-compiled for O(1) lookups during sorting + #[serde(skip)] + #[schemars(skip)] + weights: Arc>, + /// Pre-compiled glob patterns for matching with their positions + #[serde(skip)] + #[schemars(skip)] + patterns: Arc>, +} + +impl Default for ToolOrder { + fn default() -> Self { + Self { + tools: Vec::new(), + weights: Arc::new(HashMap::new()), + patterns: Arc::new(Vec::new()), + } + } +} + +impl ToolOrder { + /// Creates a new ToolOrder with the specified tool names + /// + /// # Arguments + /// + /// * `tools` - List of tool names (and patterns) to use as the basis for + /// ordering + pub fn new(tools: Vec) -> Self { + let mut weights = HashMap::new(); + let mut patterns = Vec::new(); + + for (index, tool_name) in tools.iter().enumerate() { + // Try to compile as a glob pattern + if let Ok(pattern) = Pattern::new(tool_name.as_str()) { + // Check if it's actually a pattern (contains wildcards) + if tool_name.as_str().contains('*') || tool_name.as_str().contains('?') { + patterns.push((pattern, index)); + continue; + } + } + + // Not a pattern, store as exact match + weights.insert(tool_name.clone(), index); + } + + Self { + tools, + weights: Arc::new(weights), + patterns: Arc::new(patterns), + } + } + + /// Creates a ToolOrder from a list of tool names, using the exact order + /// as specified in the list, including glob patterns. + /// + /// # Arguments + /// + /// * `tools` - List of tool names (and patterns) to use as the basis for + /// ordering + pub fn from_tool_list(tools: &[ToolName]) -> Self { + if tools.is_empty() { + return Self::default(); + } + + Self::new(tools.to_vec()) + } + + /// Sorts tool definitions according to the ordering strategy + /// + /// # Arguments + /// + /// * `tools` - Mutable slice of tool definitions to sort + pub fn sort(&self, tools: &mut [ToolDefinition]) { + if self.tools.is_empty() { + // Empty order means alphabetical + tools.sort_by(|a, b| a.name.as_str().cmp(b.name.as_str())); + } else { + tools.sort_by(|a, b| self.compare_by_weight(&a.name, &b.name)); + } + } + + /// Sorts tool definition references according to the ordering strategy + /// + /// # Arguments + /// + /// * `tools` - Mutable slice of tool definition references to sort + pub fn sort_refs(&self, tools: &mut [&ToolDefinition]) { + if self.tools.is_empty() { + // Empty order means alphabetical + tools.sort_by(|a, b| a.name.as_str().cmp(b.name.as_str())); + } else { + tools.sort_by(|a, b| self.compare_by_weight(&a.name, &b.name)); + } + } + + /// Gets the weight (position) of a tool by checking exact matches first, + /// then patterns + /// + /// Returns None if the tool doesn't match any entry in the order list. + fn get_weight(&self, tool: &ToolName) -> Option { + // First check exact match in weights map - O(1) + if let Some(&weight) = self.weights.get(tool) { + return Some(weight); + } + + // Then check glob patterns - O(p) where p is number of patterns + for (pattern, weight) in self.patterns.iter() { + if pattern.matches(tool.as_str()) { + return Some(*weight); + } + } + + None + } + + /// Compares two tool names by their weights + /// + /// Tools with lower position indices come first. Tools with no weight are + /// sorted alphabetically and appear after weighted tools. + fn compare_by_weight(&self, a: &ToolName, b: &ToolName) -> Ordering { + let a_weight = self.get_weight(a); + let b_weight = self.get_weight(b); + + match (a_weight, b_weight) { + // Both have weights - lower weight (earlier position) comes first + (Some(w_a), Some(w_b)) => match w_a.cmp(&w_b) { + // If weights are equal (e.g., both match same pattern), sort + // alphabetically + Ordering::Equal => a.as_str().cmp(b.as_str()), + other => other, + }, + // Only 'a' has weight, so it comes first + (Some(_), None) => Ordering::Less, + // Only 'b' has weight, so it comes first + (None, Some(_)) => Ordering::Greater, + // Neither has weight, sort alphabetically + (None, None) => a.as_str().cmp(b.as_str()), + } + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_alphabetical_sort() { + let fixture_order = ToolOrder::new(vec![]); // Empty list = alphabetical + let mut fixture = vec![ + ToolDefinition::new("zebra").description("Z tool"), + ToolDefinition::new("alpha").description("A tool"), + ToolDefinition::new("beta").description("B tool"), + ]; + + fixture_order.sort(&mut fixture); + + let actual: Vec = fixture.iter().map(|t| t.name.to_string()).collect(); + let expected = vec!["alpha", "beta", "zebra"]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_custom_order_all_specified() { + let fixture_order = ToolOrder::new(vec![ + ToolName::new("beta"), + ToolName::new("alpha"), + ToolName::new("zebra"), + ]); + let mut fixture = vec![ + ToolDefinition::new("zebra").description("Z tool"), + ToolDefinition::new("alpha").description("A tool"), + ToolDefinition::new("beta").description("B tool"), + ]; + + fixture_order.sort(&mut fixture); + + let actual: Vec = fixture.iter().map(|t| t.name.to_string()).collect(); + let expected = vec!["beta", "alpha", "zebra"]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_custom_order_partial_specification() { + let fixture_order = ToolOrder::new(vec![ToolName::new("zebra"), ToolName::new("beta")]); + let mut fixture = vec![ + ToolDefinition::new("alpha").description("A tool"), + ToolDefinition::new("beta").description("B tool"), + ToolDefinition::new("zebra").description("Z tool"), + ToolDefinition::new("delta").description("D tool"), + ToolDefinition::new("charlie").description("C tool"), + ]; + + fixture_order.sort(&mut fixture); + + let actual: Vec = fixture.iter().map(|t| t.name.to_string()).collect(); + // zebra and beta come first (in that order), rest alphabetically + let expected = vec!["zebra", "beta", "alpha", "charlie", "delta"]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_custom_order_with_refs() { + let fixture_order = ToolOrder::new(vec![ToolName::new("write"), ToolName::new("read")]); + let tools = [ + ToolDefinition::new("read").description("Read tool"), + ToolDefinition::new("write").description("Write tool"), + ToolDefinition::new("patch").description("Patch tool"), + ]; + let mut fixture: Vec<&ToolDefinition> = tools.iter().collect(); + + fixture_order.sort_refs(&mut fixture); + + let actual: Vec = fixture.iter().map(|t| t.name.to_string()).collect(); + let expected = vec!["write", "read", "patch"]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_custom_order_empty_list() { + let fixture_order = ToolOrder::new(vec![]); + let mut fixture = vec![ + ToolDefinition::new("zebra").description("Z tool"), + ToolDefinition::new("alpha").description("A tool"), + ToolDefinition::new("beta").description("B tool"), + ]; + + fixture_order.sort(&mut fixture); + + let actual: Vec = fixture.iter().map(|t| t.name.to_string()).collect(); + // Should fall back to alphabetical + let expected = vec!["alpha", "beta", "zebra"]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_from_tool_list_exact_order() { + let fixture = vec![ + ToolName::new("write"), + ToolName::new("read"), + ToolName::new("sage"), + ToolName::new("patch"), + ToolName::new("sem_search"), + ]; + + let actual = ToolOrder::from_tool_list(&fixture); + + let names: Vec = actual.tools.iter().map(|t| t.to_string()).collect(); + // Should maintain exact order as specified + assert_eq!(names[0], "write"); + assert_eq!(names[1], "read"); + assert_eq!(names[2], "sage"); + assert_eq!(names[3], "patch"); + assert_eq!(names[4], "sem_search"); + } + + #[test] + fn test_from_tool_list_with_mcp_tools() { + let fixture = vec![ + ToolName::new("read"), + ToolName::new("mcp_github"), + ToolName::new("write"), + ToolName::new("mcp_slack"), + ToolName::new("patch"), + ]; + + let actual = ToolOrder::from_tool_list(&fixture); + + let names: Vec = actual.tools.iter().map(|t| t.to_string()).collect(); + // Should maintain exact order as specified, no special rules + assert_eq!(names[0], "read"); + assert_eq!(names[1], "mcp_github"); + assert_eq!(names[2], "write"); + assert_eq!(names[3], "mcp_slack"); + assert_eq!(names[4], "patch"); + } + + #[test] + fn test_from_tool_list_empty() { + let fixture: Vec = vec![]; + + let actual = ToolOrder::from_tool_list(&fixture); + + assert_eq!(actual, ToolOrder::new(vec![])); + } + + #[test] + fn test_from_tool_list_with_glob_patterns() { + let fixture = vec![ + ToolName::new("read"), + ToolName::new("fs_*"), // Glob pattern - preserved + ToolName::new("write"), + ToolName::new("mcp_*"), // Glob pattern - preserved + ToolName::new("patch"), + ]; + + let actual = ToolOrder::from_tool_list(&fixture); + + let names: Vec = actual.tools.iter().map(|t| t.to_string()).collect(); + // All tools and patterns preserved + assert_eq!(names.len(), 5); + assert_eq!(names[0], "read"); + assert_eq!(names[1], "fs_*"); + assert_eq!(names[2], "write"); + assert_eq!(names[3], "mcp_*"); + assert_eq!(names[4], "patch"); + } + + #[test] + fn test_custom_order_with_glob_pattern_matching() { + let fixture_order = ToolOrder::new(vec![ + ToolName::new("read"), + ToolName::new("fs_*"), + ToolName::new("shell"), + ]); + let mut fixture = vec![ + ToolDefinition::new("shell").description("Shell tool"), + ToolDefinition::new("fs_write").description("FS Write"), + ToolDefinition::new("read").description("Read tool"), + ToolDefinition::new("fs_read").description("FS Read"), + ]; + + fixture_order.sort(&mut fixture); + + let actual: Vec = fixture.iter().map(|t| t.name.to_string()).collect(); + // read (pos 0), fs_read and fs_write (both match fs_* at pos 1, alphabetically + // sorted), shell (pos 2) + let expected = vec!["read", "fs_read", "fs_write", "shell"]; + + assert_eq!(actual, expected); + } + + #[test] + fn test_is_empty() { + let empty = ToolOrder::new(vec![]); + let non_empty = ToolOrder::new(vec![ToolName::new("read")]); + + assert!(empty.tools.is_empty()); + assert!(!non_empty.tools.is_empty()); + } + + #[test] + fn test_weight_lookup_optimization() { + // Test that the HashMap optimization works correctly + let fixture_order = ToolOrder::new(vec![ + ToolName::new("alpha"), + ToolName::new("beta"), + ToolName::new("gamma"), + ]); + + // Verify weights are stored correctly + assert_eq!(fixture_order.get_weight(&ToolName::new("alpha")), Some(0)); + assert_eq!(fixture_order.get_weight(&ToolName::new("beta")), Some(1)); + assert_eq!(fixture_order.get_weight(&ToolName::new("gamma")), Some(2)); + assert_eq!(fixture_order.get_weight(&ToolName::new("delta")), None); + } + + #[test] + fn test_pattern_matching_optimization() { + // Test that patterns are pre-compiled + let fixture_order = ToolOrder::new(vec![ + ToolName::new("exact_match"), + ToolName::new("prefix_*"), + ToolName::new("*_suffix"), + ]); + + // Exact matches should be in weights + assert_eq!( + fixture_order.get_weight(&ToolName::new("exact_match")), + Some(0) + ); + + // Pattern matches should work + assert_eq!( + fixture_order.get_weight(&ToolName::new("prefix_test")), + Some(1) + ); + assert_eq!( + fixture_order.get_weight(&ToolName::new("prefix_foo")), + Some(1) + ); + assert_eq!( + fixture_order.get_weight(&ToolName::new("test_suffix")), + Some(2) + ); + + // Non-matches should return None + assert_eq!(fixture_order.get_weight(&ToolName::new("no_match")), None); + } + + #[test] + fn test_clone_is_cheap() { + // Test that cloning ToolOrder is cheap (Arc is cloned, not the data) + let fixture_order = ToolOrder::new(vec![ + ToolName::new("alpha"), + ToolName::new("beta"), + ToolName::new("gamma"), + ]); + + let cloned = fixture_order.clone(); + + // Both should have the same Arc pointers (cheap clone) + assert_eq!( + Arc::as_ptr(&fixture_order.weights), + Arc::as_ptr(&cloned.weights) + ); + assert_eq!( + Arc::as_ptr(&fixture_order.patterns), + Arc::as_ptr(&cloned.patterns) + ); + } +} diff --git a/crates/forge_domain/src/top_k.rs b/crates/forge_domain/src/top_k.rs new file mode 100644 index 0000000000000000000000000000000000000000..daa2d385d429f6be34aaa2abaefae15a6564b53d --- /dev/null +++ b/crates/forge_domain/src/top_k.rs @@ -0,0 +1,194 @@ +use std::fmt; +use std::ops::Deref; + +use schemars::JsonSchema; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +/// A newtype for top_k values with built-in validation +/// +/// Top-k controls the number of highest probability vocabulary tokens to keep: +/// - Lower values (e.g., 10) make responses more focused by considering only +/// the top K most likely tokens +/// - Higher values (e.g., 100) make responses more diverse by considering more +/// token options +/// - Valid range is 1 to 1000 (inclusive) +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, JsonSchema)] +pub struct TopK(u32); + +impl TopK { + /// Creates a new TopK value, returning an error if outside the valid + /// range (1 to 1000) + pub fn new(value: u32) -> Result { + if Self::is_valid(value) { + Ok(Self(value)) + } else { + Err(format!("top_k must be between 1 and 1000, got {value}")) + } + } + + /// Creates a new TopK value without validation + /// + /// # Safety + /// This function should only be used when the value is known to be valid + pub fn new_unchecked(value: u32) -> Self { + debug_assert!(Self::is_valid(value), "invalid top_k: {value}"); + Self(value) + } + + /// Returns true if the top_k value is within the valid range (1 to 1000) + pub fn is_valid(value: u32) -> bool { + (1..=1000).contains(&value) + } + + /// Returns the inner u32 value + pub fn value(&self) -> u32 { + self.0 + } +} + +impl Deref for TopK { + type Target = u32; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl From for u32 { + fn from(top_k: TopK) -> Self { + top_k.0 + } +} + +impl fmt::Display for TopK { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +impl Serialize for TopK { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_u32(self.0) + } +} + +impl<'de> Deserialize<'de> for TopK { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + use serde::de::Error; + let value = u32::deserialize(deserializer)?; + if Self::is_valid(value) { + Ok(Self(value)) + } else { + Err(Error::custom(format!( + "top_k must be between 1 and 1000, got {value}" + ))) + } + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use serde_json::json; + + use super::*; + + #[test] + fn test_top_k_creation() { + // Valid top_k values should be created successfully + let valid_values = [1, 10, 50, 100, 500, 1000]; + for value in valid_values { + let result = TopK::new(value); + assert!(result.is_ok(), "TopK {value} should be valid"); + assert_eq!(result.unwrap().value(), value); + } + + // Invalid top_k values should return an error + let invalid_values = [0, 1001, 2000, 5000]; + for value in invalid_values { + let result = TopK::new(value); + assert!(result.is_err(), "TopK {value} should be invalid"); + assert!( + result + .unwrap_err() + .contains("top_k must be between 1 and 1000"), + "Error should mention valid range" + ); + } + } + + #[test] + fn test_top_k_serialization() { + let top_k = TopK::new(50).unwrap(); + let json = serde_json::to_value(top_k).unwrap(); + + if let serde_json::Value::Number(num) = &json { + let int_val = num.as_u64().unwrap(); + assert_eq!(int_val, 50); + } else { + panic!("Expected a number, got {json:?}"); + } + } + + #[test] + fn test_top_k_deserialization() { + // Valid top_k values should deserialize correctly + let valid_values = [1, 10, 50, 100, 500, 1000]; + for value in valid_values { + let json = json!(value); + let top_k: Result = serde_json::from_value(json); + assert!(top_k.is_ok(), "Valid top_k {value} should deserialize"); + assert_eq!(top_k.unwrap().value(), value); + } + + // Invalid top_k values should fail deserialization + let invalid_values = [0, 1001, 2000, 5000]; + for value in invalid_values { + let json = json!(value); + let top_k: Result = serde_json::from_value(json); + assert!( + top_k.is_err(), + "Invalid top_k {value} should fail deserialization" + ); + let err = top_k.unwrap_err().to_string(); + assert!( + err.contains("top_k must be between 1 and 1000"), + "Error should mention valid range: {err}" + ); + } + } + + #[test] + fn test_top_k_in_struct() { + #[derive(Serialize, Deserialize, Debug)] + struct TestStruct { + top_k: TopK, + } + + // Valid top_k + let json = json!({ + "top_k": 50 + }); + let test_struct: Result = serde_json::from_value(json); + assert!(test_struct.is_ok()); + assert_eq!(test_struct.unwrap().top_k.value(), 50); + + // Invalid top_k + let json = json!({ + "top_k": 1500 + }); + let test_struct: Result = serde_json::from_value(json); + assert!(test_struct.is_err()); + let err = test_struct.unwrap_err().to_string(); + assert!( + err.contains("top_k must be between 1 and 1000"), + "Error should mention valid range: {err}" + ); + } +} diff --git a/crates/forge_domain/src/update.rs b/crates/forge_domain/src/update.rs new file mode 100644 index 0000000000000000000000000000000000000000..1838d67524af1bab643b1641de2996d39c79c2d2 --- /dev/null +++ b/crates/forge_domain/src/update.rs @@ -0,0 +1,34 @@ +use std::time::Duration; + +use derive_setters::Setters; +use merge::Merge; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +#[derive(Default, Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum UpdateFrequency { + Daily, + Weekly, + Never, + #[default] + Always, +} + +impl From for Duration { + fn from(val: UpdateFrequency) -> Self { + match val { + UpdateFrequency::Daily => Duration::from_secs(60 * 60 * 24), // 1 day + UpdateFrequency::Weekly => Duration::from_secs(60 * 60 * 24 * 7), // 1 week + UpdateFrequency::Never => Duration::MAX, + UpdateFrequency::Always => Duration::ZERO, // one time, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Merge, Default, JsonSchema, Setters, PartialEq)] +#[merge(strategy = merge::option::overwrite_none)] +pub struct Update { + pub frequency: Option, + pub auto_update: Option, +} diff --git a/crates/forge_domain/src/validation.rs b/crates/forge_domain/src/validation.rs new file mode 100644 index 0000000000000000000000000000000000000000..bc26e5f7843821236de4d354d83695fee0bcc92d --- /dev/null +++ b/crates/forge_domain/src/validation.rs @@ -0,0 +1,10 @@ +/// Represents a single syntax error in a file +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SyntaxError { + /// Line number where the error occurred (1-based) + pub line: u32, + /// Column number where the error occurred (1-based) + pub column: u32, + /// Error message describing the syntax issue + pub message: String, +} diff --git a/crates/forge_domain/src/workspace.rs b/crates/forge_domain/src/workspace.rs new file mode 100644 index 0000000000000000000000000000000000000000..3e88b44016e465786e5be969aa652fb3f33be993 --- /dev/null +++ b/crates/forge_domain/src/workspace.rs @@ -0,0 +1,30 @@ +use derive_more::Display; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +/// Workspace identifier (UUID) from workspace server. +/// +/// Generated locally and sent to server during CreateWorkspace. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Display)] +#[display("{}", _0)] +pub struct WorkspaceId(Uuid); + +impl WorkspaceId { + /// Generate a new random workspace ID + pub fn generate() -> Self { + Self(Uuid::new_v4()) + } + + /// Parse a workspace ID from a string + /// + /// # Errors + /// Returns an error if the string is not a valid UUID + pub fn from_string(s: &str) -> anyhow::Result { + Ok(Self(Uuid::parse_str(s)?)) + } + + /// Get the inner UUID + pub fn inner(&self) -> Uuid { + self.0 + } +} diff --git a/crates/forge_domain/src/xml.rs b/crates/forge_domain/src/xml.rs new file mode 100644 index 0000000000000000000000000000000000000000..7ba5fbadfcc061d5d15acc40d03874b51a444781 --- /dev/null +++ b/crates/forge_domain/src/xml.rs @@ -0,0 +1,168 @@ +/// Extracts content between the specified XML-style tags +/// +/// # Arguments +/// +/// * `text` - The text to extract content from +/// * `tag_name` - The name of the XML tag (without angle brackets) +/// +/// # Returns +/// +/// * `Some(&str)` containing the extracted content if tags are found +/// * `None` if the tags are not found +pub fn extract_tag_content<'a>(text: &'a str, tag_name: &str) -> Option<&'a str> { + let opening_tag = format!("<{tag_name}>",); + let closing_tag = format!(""); + + #[allow(clippy::collapsible_if)] + if let Some(start_idx) = text.find(&opening_tag) { + if let Some(end_idx) = text.rfind(&closing_tag) { + let content_start = start_idx + opening_tag.len(); + if content_start < end_idx { + return text.get(content_start..end_idx).map(|s| s.trim()); + } + } + } + + None +} + +/// Removes content within XML-style tags that start with the specified prefix +pub fn remove_tag_with_prefix(text: &str, prefix: &str) -> String { + // First, find all unique tag names that start with the prefix + let tag_pattern = format!(r"<({prefix}[a-zA-Z0-9_-]*?)(?:\s[^>]*?)?>"); + let mut tag_names = Vec::new(); + + if let Ok(regex) = regex::Regex::new(&tag_pattern) { + for captures in regex.captures_iter(text) { + if let Some(tag_name) = captures.get(1) { + // Only add unique tag names to the list + let tag_name = tag_name.as_str().to_string(); + if !tag_names.contains(&tag_name) { + tag_names.push(tag_name); + } + } + } + } + + // Now remove content for each tag name found + let mut result = text.to_string(); + for tag_name in tag_names { + // Create pattern to match complete tag including content + let pattern = format!(r"<{tag_name}(?:\s[^>]*?)?>[\s\S]*?"); + + if let Ok(regex) = regex::Regex::new(&pattern) { + result = regex.replace_all(&result, "").to_string(); + } + } + + result +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_extract_tag_content() { + let fixture = "Some text

This is the important part and more text"; + let actual = extract_tag_content(fixture, "summary"); + let expected = Some("This is the important part"); + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_tag_content_no_tags() { + let fixture = "Some text without any tags"; + let actual = extract_tag_content(fixture, "summary"); + let expected = None; + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_tag_content_with_different_tag() { + let fixture = "Text with Custom content tags"; + let actual = extract_tag_content(fixture, "custom"); + let expected = Some("Custom content"); + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_tag_content_with_malformed_tags() { + let fixture = "Text with but no closing tag"; + let actual = extract_tag_content(fixture, "opening"); + let expected = None; + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_tag_names_with_prefix() { + let fixture = "Something Content More"; + let actual = remove_tag_with_prefix(fixture, "forge"); + // Check that both tool tags have been removed, leaving only tag + assert!(actual.contains("More")); + assert!(!actual.contains("")); + assert!(!actual.contains("")); + } + + #[test] + fn test_extract_tag_names_with_prefix_no_matches() { + let fixture = "Some content Other content"; + let actual = remove_tag_with_prefix(fixture, "forge"); + let expected = "Some content Other content"; + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_tag_names_with_prefix_nested() { + let fixture = "InnerNested"; + let actual = remove_tag_with_prefix(fixture, "forge"); + let expected = ""; + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_tag_names_with_prefix_duplicates() { + let fixture = + "FirstMiddleSecond"; + let actual = remove_tag_with_prefix(fixture, "forge"); + let expected = "Middle"; + assert_eq!(actual, expected); + } + + #[test] + fn test_extract_tag_names_with_prefix_attributes() { + let fixture = "Content More"; + let actual = remove_tag_with_prefix(fixture, "forge"); + // Check that both tool tags have been removed + assert!(!actual.contains(" Some(text), + _ => None, + }) + .expect("Should have assistant message"); + + let tool_calls = assistant_msg + .tool_calls + .as_ref() + .expect("Should have tool calls"); + assert_eq!(tool_calls.len(), 1); + + let tool_call = &tool_calls[0]; + assert_eq!(tool_call.name.as_str(), "read"); + + // Verify arguments are parsed correctly (can access fields) + let parsed_args = tool_call.arguments.parse().expect("Should parse arguments"); + assert_eq!(parsed_args["file_path"], "/test/path"); + assert_eq!(parsed_args["range"]["start_line"], 1); + + // Now serialize the context back to JSON (as we would send to API) + let serialized = serde_json::to_string(&context).expect("Should serialize"); + println!("Serialized: {}", serialized); + + // Parse the JSON to verify structure + let reparsed: serde_json::Value = serde_json::from_str(&serialized).expect("Should re-parse"); + + // Find the assistant message + let messages = reparsed["messages"] + .as_array() + .expect("Should have messages"); + let assistant_json = messages + .iter() + .find(|m| m["text"]["role"] == "Assistant") + .expect("Should find assistant message"); + + let tool_calls_json = assistant_json["text"]["tool_calls"] + .as_array() + .expect("Should have tool_calls"); + + let args = &tool_calls_json[0]["arguments"]; + + // THE KEY TEST: arguments should be a JSON object, not a string + // This is what Fireworks API expects - if it's a string, we get 400 error + assert!( + args.is_object(), + "CRITICAL: arguments must be a JSON object for API, not a string. Got: {}", + args + ); + + // Verify the values are preserved correctly + assert_eq!(args["file_path"], "/test/path"); + assert_eq!(args["range"]["start_line"], 1); + assert_eq!(args["range"]["end_line"], 10); + + println!("SUCCESS: Stringified arguments properly converted to JSON object"); +} + +/// Test with patch tool (the actual error case from kimi-k2p5-turbo) +#[test] +fn test_kimi_k2p5_turbo_patch_tool_scenario() { + // This simulates the exact error case: patch tool with stringified arguments + let conversation_json = r#"{ + "messages": [ + { + "text": { + "role": "System", + "content": "You are Forge." + } + }, + { + "text": { + "role": "User", + "content": "Edit the file" + } + }, + { + "text": { + "role": "Assistant", + "content": "I'll edit the file.", + "tool_calls": [ + { + "name": "patch", + "call_id": "call_patch_001", + "arguments": "{\"file_path\": \"/test/file.rs\", \"old_string\": \"fn main() {\\n}\", \"new_string\": \"fn main() {\\n println!(\\\"Hello\\\");\\n}\"}" + } + ] + } + } + ] + }"#; + + let context: Context = + serde_json::from_str(conversation_json).expect("Failed to parse conversation"); + + // Serialize back to what we'd send to API + let serialized = serde_json::to_string(&context).expect("Should serialize"); + let reparsed: serde_json::Value = serde_json::from_str(&serialized).expect("Should parse"); + + let messages = reparsed["messages"] + .as_array() + .expect("Should have messages"); + let assistant = messages + .iter() + .find(|m| m["text"]["role"] == "Assistant") + .expect("Should find assistant"); + + let tool_calls = assistant["text"]["tool_calls"] + .as_array() + .expect("Should have tool_calls"); + let args = &tool_calls[0]["arguments"]; + + // THE CRITICAL ASSERTION for the kimi-k2p5-turbo fix + assert!( + args.is_object(), + "Tool call arguments must be a JSON object when sent to API, not a string. Got: {}", + args + ); + + // Verify nested strings with special characters are preserved + assert_eq!(args["file_path"], "/test/file.rs"); + assert_eq!(args["old_string"], "fn main() {\n}"); + assert_eq!( + args["new_string"], + "fn main() {\n println!(\"Hello\");\n}" + ); + + println!("SUCCESS: Patch tool with stringified args properly converted"); +} + +/// Test multiple tool calls with stringified arguments +#[test] +fn test_multiple_stringified_tool_calls() { + let conversation_json = r#"{ + "messages": [ + { + "text": { + "role": "System", + "content": "You are Forge." + } + }, + { + "text": { + "role": "User", + "content": "Do multiple things" + } + }, + { + "text": { + "role": "Assistant", + "content": "I'll do multiple things.", + "tool_calls": [ + { + "name": "read", + "call_id": "call_1", + "arguments": "{\"file_path\": \"/file1.txt\"}" + }, + { + "name": "patch", + "call_id": "call_2", + "arguments": "{\"file_path\": \"/file2.txt\", \"old_string\": \"a\", \"new_string\": \"b\"}" + } + ] + } + } + ] + }"#; + + let context: Context = + serde_json::from_str(conversation_json).expect("Failed to parse conversation"); + let serialized = serde_json::to_string(&context).expect("Should serialize"); + let reparsed: serde_json::Value = serde_json::from_str(&serialized).expect("Should parse"); + + let messages = reparsed["messages"] + .as_array() + .expect("Should have messages"); + let assistant = messages + .iter() + .find(|m| m["text"]["role"] == "Assistant") + .expect("Should find assistant"); + + let tool_calls = assistant["text"]["tool_calls"] + .as_array() + .expect("Should have tool_calls"); + + assert_eq!(tool_calls.len(), 2); + + // Both tool calls should have arguments as objects, not strings + for (i, tc) in tool_calls.iter().enumerate() { + let args = &tc["arguments"]; + assert!( + args.is_object(), + "Tool call {} arguments should be a JSON object, got: {}", + i, + args + ); + } + + println!("SUCCESS: Multiple stringified tool calls properly converted to JSON objects"); +} + +/// Test that regular JSON objects (not stringified) still work correctly +#[test] +fn test_regular_json_objects_unchanged() { + // Normal case: arguments are already proper JSON objects + let conversation_json = r#"{ + "messages": [ + { + "text": { + "role": "System", + "content": "You are Forge." + } + }, + { + "text": { + "role": "Assistant", + "content": "I'll read the file.", + "tool_calls": [ + { + "name": "read", + "call_id": "call_001", + "arguments": {"file_path": "/test/path", "range": {"start_line": 1, "end_line": 10}} + } + ] + } + } + ] + }"#; + + let context: Context = + serde_json::from_str(conversation_json).expect("Failed to parse conversation"); + let serialized = serde_json::to_string(&context).expect("Should serialize"); + let reparsed: serde_json::Value = serde_json::from_str(&serialized).expect("Should parse"); + + let messages = reparsed["messages"] + .as_array() + .expect("Should have messages"); + let assistant = messages + .iter() + .find(|m| m["text"]["role"] == "Assistant") + .expect("Should find assistant"); + + let tool_calls = assistant["text"]["tool_calls"] + .as_array() + .expect("Should have tool_calls"); + let args = &tool_calls[0]["arguments"]; + + // Should still be an object + assert!( + args.is_object(), + "Regular JSON objects should remain as objects" + ); + assert_eq!(args["file_path"], "/test/path"); + assert_eq!(args["range"]["start_line"], 1); + + println!("SUCCESS: Regular JSON objects work correctly"); +} diff --git a/crates/forge_eventsource_stream/src/event.rs b/crates/forge_eventsource_stream/src/event.rs new file mode 100644 index 0000000000000000000000000000000000000000..651b5afadf524967beb2823c30f469cfae99d7de --- /dev/null +++ b/crates/forge_eventsource_stream/src/event.rs @@ -0,0 +1,16 @@ +#[cfg(not(feature = "std"))] +use alloc::string::String; +use core::time::Duration; + +/// An Event +#[derive(Default, Debug, Eq, PartialEq, Clone)] +pub struct Event { + /// The event name if given + pub event: String, + /// The event data + pub data: String, + /// The event id if given + pub id: String, + /// Retry duration if given + pub retry: Option, +} diff --git a/crates/forge_eventsource_stream/src/event_stream.rs b/crates/forge_eventsource_stream/src/event_stream.rs new file mode 100644 index 0000000000000000000000000000000000000000..cbdc4a31013145ac3e3953774d18638e3d51d3be --- /dev/null +++ b/crates/forge_eventsource_stream/src/event_stream.rs @@ -0,0 +1,560 @@ +#[cfg(not(feature = "std"))] +use alloc::string::{FromUtf8Error, String, ToString}; +use core::fmt; +use core::pin::Pin; +use core::time::Duration; +#[cfg(feature = "std")] +use std::string::FromUtf8Error; + +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use nom::error::Error as NomError; +use pin_project_lite::pin_project; + +use crate::event::Event; +use crate::parser::{RawEventLine, is_bom, is_lf, line}; +use crate::utf8_stream::{Utf8Stream, Utf8StreamError}; + +#[derive(Default, Debug)] +struct EventBuilder { + event: Event, + is_complete: bool, +} + +impl EventBuilder { + /// From the HTML spec + /// + /// -> If the field name is "event" + /// Set the event type buffer to field value. + /// + /// -> If the field name is "data" + /// Append the field value to the data buffer, then append a single + /// U+000A LINE FEED (LF) character to the data buffer. + /// + /// -> If the field name is "id" + /// If the field value does not contain U+0000 NULL, then set the last + /// event ID buffer to the field value. Otherwise, ignore the field. + /// + /// -> If the field name is "retry" + /// If the field value consists of only ASCII digits, then interpret the + /// field value as an integer in base ten, and set the event stream's + /// reconnection time to that integer. Otherwise, ignore the field. + /// + /// -> Otherwise + /// The field is ignored. + fn add(&mut self, line: RawEventLine) { + match line { + RawEventLine::Field(field, val) => { + let val = val.unwrap_or(""); + match field { + "event" => { + self.event.event = val.to_string(); + } + "data" => { + self.event.data.push_str(val); + self.event.data.push('\u{000A}'); + } + "id" if !val.contains('\u{0000}') => self.event.id = val.to_string(), + "retry" => { + if let Ok(val) = val.parse::() { + self.event.retry = Some(Duration::from_millis(val)) + } + } + _ => {} + } + } + RawEventLine::Comment(_) => {} + RawEventLine::Empty => self.is_complete = true, + } + } + + /// From the HTML spec + /// + /// 1. Set the last event ID string of the event source to the value of the + /// last event ID buffer. The buffer does not get reset, so the last + /// event ID string of the event source remains set to this value until + /// the next time it is set by the server. + /// 2. If the data buffer is an empty string, set the data buffer and the + /// event type buffer to the empty string and return. + /// 3. If the data buffer's last character is a U+000A LINE FEED (LF) + /// character, then remove the last character from the data buffer. + /// 4. Let event be the result of creating an event using MessageEvent, in + /// the relevant Realm of the EventSource object. + /// 5. Initialize event's type attribute to message, its data attribute to + /// data, its origin attribute to the serialization of the origin of the + /// event stream's final URL (i.e., the URL after redirects), and its + /// lastEventId attribute to the last event ID string of the event + /// source. + /// 6. If the event type buffer has a value other than the empty string, + /// change the type of the newly created event to equal the value of the + /// event type buffer. + /// 7. Set the data buffer and the event type buffer to the empty string. + /// 8. Queue a task which, if the readyState attribute is set to a value + /// other than CLOSED, dispatches the newly created event at the + /// EventSource object. + fn dispatch(&mut self) -> Option { + let builder = core::mem::take(self); + let mut event = builder.event; + self.event.id = event.id.clone(); + + if event.data.is_empty() { + return None; + } + + if is_lf(event.data.chars().next_back().unwrap()) { + event.data.pop(); + } + + if event.event.is_empty() { + event.event = "message".to_string(); + } + + Some(event) + } +} + +#[derive(Debug, Clone, Copy)] +pub enum EventStreamState { + NotStarted, + Started, + Terminated, +} + +impl EventStreamState { + fn is_terminated(self) -> bool { + matches!(self, Self::Terminated) + } + fn is_started(self) -> bool { + matches!(self, Self::Started) + } +} + +pin_project! { +/// A Stream of events +pub struct EventStream { + #[pin] + stream: Utf8Stream, + buffer: String, + builder: EventBuilder, + state: EventStreamState, + last_event_id: String, +} +} + +impl EventStream { + /// Initialize the EventStream with a Stream + pub fn new(stream: S) -> Self { + Self { + stream: Utf8Stream::new(stream), + buffer: String::new(), + builder: EventBuilder::default(), + state: EventStreamState::NotStarted, + last_event_id: String::new(), + } + } + + /// Set the last event ID of the stream. Useful for initializing the stream + /// with a previous last event ID + pub fn set_last_event_id(&mut self, id: impl Into) { + self.last_event_id = id.into(); + } + + /// Get the last event ID of the stream + pub fn last_event_id(&self) -> &str { + &self.last_event_id + } +} + +/// Error thrown while parsing an event line +#[derive(Debug, PartialEq)] +pub enum EventStreamError { + /// Source stream is not valid UTF8 + Utf8(FromUtf8Error), + /// Source stream is not a valid EventStream + Parser(NomError), + /// Underlying source stream error + Transport(E), +} + +impl From> for EventStreamError { + fn from(err: Utf8StreamError) -> Self { + match err { + Utf8StreamError::Utf8(err) => Self::Utf8(err), + Utf8StreamError::Transport(err) => Self::Transport(err), + } + } +} + +impl From> for EventStreamError { + fn from(err: NomError<&str>) -> Self { + EventStreamError::Parser(NomError::new(err.input.to_string(), err.code)) + } +} + +impl fmt::Display for EventStreamError +where + E: fmt::Display, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Utf8(err) => f.write_fmt(format_args!("UTF8 error: {}", err)), + Self::Parser(err) => f.write_fmt(format_args!("Parse error: {}", err)), + Self::Transport(err) => f.write_fmt(format_args!("Transport error: {}", err)), + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for EventStreamError where E: fmt::Display + fmt::Debug + Send + Sync {} + +fn parse_event( + buffer: &mut String, + builder: &mut EventBuilder, +) -> Result, EventStreamError> { + if buffer.is_empty() { + return Ok(None); + } + loop { + match line(buffer.as_ref()) { + Ok((rem, next_line)) => { + builder.add(next_line); + let consumed = buffer.len() - rem.len(); + let rem = buffer.split_off(consumed); + *buffer = rem; + if builder.is_complete + && let Some(event) = builder.dispatch() + { + return Ok(Some(event)); + } + } + Err(nom::Err::Incomplete(_)) => return Ok(None), + Err(nom::Err::Error(err)) | Err(nom::Err::Failure(err)) => return Err(err.into()), + } + } +} + +impl Stream for EventStream +where + S: Stream>, + B: AsRef<[u8]>, +{ + type Item = Result>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll> { + let mut this = self.project(); + + match parse_event(this.buffer, this.builder) { + Ok(Some(event)) => { + *this.last_event_id = event.id.clone(); + return Poll::Ready(Some(Ok(event))); + } + Err(err) => return Poll::Ready(Some(Err(err))), + _ => {} + } + + if this.state.is_terminated() { + return Poll::Ready(None); + } + + loop { + match this.stream.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(string))) => { + if string.is_empty() { + continue; + } + + let slice = if this.state.is_started() { + string.as_str() + } else { + *this.state = EventStreamState::Started; + string.strip_prefix(is_bom).unwrap_or(&string) + }; + this.buffer.push_str(slice); + + match parse_event(this.buffer, this.builder) { + Ok(Some(event)) => { + *this.last_event_id = event.id.clone(); + return Poll::Ready(Some(Ok(event))); + } + Err(err) => return Poll::Ready(Some(Err(err))), + _ => {} + } + } + Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err.into()))), + Poll::Ready(None) => { + *this.state = EventStreamState::Terminated; + return Poll::Ready(None); + } + Poll::Pending => return Poll::Pending, + } + } + } +} + +#[cfg(test)] +mod tests { + use futures::prelude::*; + + use super::*; + + #[tokio::test] + async fn valid_data_fields() { + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + "data: Hello, world!\n\n" + )])) + .try_collect::>() + .await + .unwrap(), + vec![Event { + event: "message".to_string(), + data: "Hello, world!".to_string(), + ..Default::default() + }] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![ + Ok::<_, ()>("data: Hello,"), + Ok::<_, ()>(" world!\n\n") + ])) + .try_collect::>() + .await + .unwrap(), + vec![Event { + event: "message".to_string(), + data: "Hello, world!".to_string(), + ..Default::default() + }] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![ + Ok::<_, ()>("data: Hello,"), + Ok::<_, ()>(""), + Ok::<_, ()>(" world!\n\n") + ])) + .try_collect::>() + .await + .unwrap(), + vec![Event { + event: "message".to_string(), + data: "Hello, world!".to_string(), + ..Default::default() + }] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + "data: Hello, world!\n" + )])) + .try_collect::>() + .await + .unwrap(), + vec![] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + "data: Hello,\ndata: world!\n\n" + )])) + .try_collect::>() + .await + .unwrap(), + vec![Event { + event: "message".to_string(), + data: "Hello,\nworld!".to_string(), + ..Default::default() + }] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + "data: Hello,\n\ndata: world!\n\n" + )])) + .try_collect::>() + .await + .unwrap(), + vec![ + Event { + event: "message".to_string(), + data: "Hello,".to_string(), + ..Default::default() + }, + Event { + event: "message".to_string(), + data: "world!".to_string(), + ..Default::default() + } + ] + ); + } + + #[tokio::test] + async fn spec_examples() { + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + "data: This is the first message. + +data: This is the second message, it +data: has two lines. + +data: This is the third message. + +" + )])) + .try_collect::>() + .await + .unwrap(), + vec![ + Event { + event: "message".to_string(), + data: "This is the first message.".to_string(), + ..Default::default() + }, + Event { + event: "message".to_string(), + data: "This is the second message, it\nhas two lines.".to_string(), + ..Default::default() + }, + Event { + event: "message".to_string(), + data: "This is the third message.".to_string(), + ..Default::default() + } + ] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + "event: add +data: 73857293 + +event: remove +data: 2153 + +event: add +data: 113411 + +" + )])) + .try_collect::>() + .await + .unwrap(), + vec![ + Event { + event: "add".to_string(), + data: "73857293".to_string(), + ..Default::default() + }, + Event { + event: "remove".to_string(), + data: "2153".to_string(), + ..Default::default() + }, + Event { + event: "add".to_string(), + data: "113411".to_string(), + ..Default::default() + } + ] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + "data: YHOO +data: +2 +data: 10 + +" + )])) + .try_collect::>() + .await + .unwrap(), + vec![Event { + event: "message".to_string(), + data: "YHOO\n+2\n10".to_string(), + ..Default::default() + },] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + ": test stream + +data: first event +id: 1 + +data:second event +id + +data: third event + +" + )])) + .try_collect::>() + .await + .unwrap(), + vec![ + Event { + event: "message".to_string(), + id: "1".to_string(), + data: "first event".to_string(), + ..Default::default() + }, + Event { + event: "message".to_string(), + data: "second event".to_string(), + ..Default::default() + }, + Event { + event: "message".to_string(), + data: " third event".to_string(), + ..Default::default() + } + ] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + "data + +data +data + +data: +" + )])) + .try_collect::>() + .await + .unwrap(), + vec![ + Event { + event: "message".to_string(), + data: "".to_string(), + ..Default::default() + }, + Event { + event: "message".to_string(), + data: "\n".to_string(), + ..Default::default() + }, + ] + ); + assert_eq!( + EventStream::new(futures::stream::iter(vec![Ok::<_, ()>( + "data:test + +data: test + +" + )])) + .try_collect::>() + .await + .unwrap(), + vec![ + Event { + event: "message".to_string(), + data: "test".to_string(), + ..Default::default() + }, + Event { + event: "message".to_string(), + data: "test".to_string(), + ..Default::default() + }, + ] + ); + } +} diff --git a/crates/forge_eventsource_stream/src/lib.rs b/crates/forge_eventsource_stream/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..0368336aabeaf1bfb3627c1c4712a7bc7b9bad3e --- /dev/null +++ b/crates/forge_eventsource_stream/src/lib.rs @@ -0,0 +1,41 @@ +//! A basic building block for building an Eventsource from a Stream of bytes +//! array like objects. To learn more about Server Sent Events (SSE) take a look +//! at [the MDN docs](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events) +//! +//! # Example +//! +//! ```ignore +//! let mut stream = reqwest::Client::new() +//! .get("http://localhost:7020/notifications") +//! .send() +//! .await? +//! .bytes_stream() +//! .eventsource(); +//! +//! +//! while let Some(event) = stream.next().await { +//! match event { +//! Ok(event) => println!( +//! "received event[type={}]: {}", +//! event.event, +//! event.data +//! ), +//! Err(e) => eprintln!("error occured: {}", e), +//! } +//! } +//! ``` + +#![cfg_attr(not(feature = "std"), no_std)] + +#[cfg(not(feature = "std"))] +extern crate alloc; + +mod event; +mod event_stream; +mod parser; +mod traits; +mod utf8_stream; + +pub use event::Event; +pub use event_stream::{EventStream, EventStreamError}; +pub use traits::Eventsource; diff --git a/crates/forge_eventsource_stream/src/parser.rs b/crates/forge_eventsource_stream/src/parser.rs new file mode 100644 index 0000000000000000000000000000000000000000..3e8a279b8c84c2b4c6bb076c8344709d41362a06 --- /dev/null +++ b/crates/forge_eventsource_stream/src/parser.rs @@ -0,0 +1,124 @@ +use nom::branch::alt; +use nom::bytes::complete::take_while_m_n as take_while_m_n_complete; +use nom::bytes::streaming::{tag, take_while, take_while_m_n, take_while1}; +use nom::combinator::opt; +use nom::sequence::{preceded, terminated}; +use nom::{IResult, Parser}; + +/// ; ABNF definition from HTML spec +/// +/// stream = [ bom ] *event +/// event = *( comment / field ) end-of-line +/// comment = colon *any-char end-of-line +/// field = 1*name-char [ colon [ space ] *any-char ] end-of-line +/// end-of-line = ( cr lf / cr / lf ) +/// +/// ; characters +/// lf = %x000A ; U+000A LINE FEED (LF) +/// cr = %x000D ; U+000D CARRIAGE RETURN (CR) +/// space = %x0020 ; U+0020 SPACE +/// colon = %x003A ; U+003A COLON (:) +/// bom = %xFEFF ; U+FEFF BYTE ORDER MARK +/// name-char = %x0000-0009 / %x000B-000C / %x000E-0039 / %x003B-10FFFF +/// ; a scalar value other than U+000A LINE FEED (LF), U+000D +/// CARRIAGE RETURN (CR), or U+003A COLON (:) any-char = %x0000-0009 / +/// %x000B-000C / %x000E-10FFFF ; a scalar value other than +/// U+000A LINE FEED (LF) or U+000D CARRIAGE RETURN (CR) + +#[derive(Debug)] +#[allow(dead_code)] +pub enum RawEventLine<'a> { + Comment(&'a str), + Field(&'a str, Option<&'a str>), + Empty, +} + +#[inline] +pub fn is_lf(c: char) -> bool { + c == '\u{000A}' +} + +#[inline] +pub fn is_cr(c: char) -> bool { + c == '\u{000D}' +} + +#[inline] +pub fn is_space(c: char) -> bool { + c == '\u{0020}' +} + +#[inline] +pub fn is_colon(c: char) -> bool { + c == '\u{003A}' +} + +#[inline] +pub fn is_bom(c: char) -> bool { + c == '\u{feff}' +} + +#[inline] +pub fn is_name_char(c: char) -> bool { + matches!(c, '\u{0000}'..='\u{0009}' + | '\u{000B}'..='\u{000C}' + | '\u{000E}'..='\u{0039}' + | '\u{003B}'..='\u{10FFFF}') +} + +#[inline] +pub fn is_any_char(c: char) -> bool { + matches!(c, '\u{0000}'..='\u{0009}' + | '\u{000B}'..='\u{000C}' + | '\u{000E}'..='\u{10FFFF}') +} + +#[inline] +fn crlf(input: &str) -> IResult<&str, &str> { + tag("\u{000D}\u{000A}").parse(input) +} + +#[inline] +fn end_of_line(input: &str) -> IResult<&str, &str> { + alt(( + crlf, + take_while_m_n_complete(1, 1, is_cr), + take_while_m_n_complete(1, 1, is_lf), + )) + .parse(input) +} + +#[inline] +fn comment(input: &str) -> IResult<&str, RawEventLine<'_>> { + preceded( + take_while_m_n(1, 1, is_colon), + terminated(take_while(is_any_char), end_of_line), + ) + .parse(input) + .map(|(input, comment)| (input, RawEventLine::Comment(comment))) +} + +#[inline] +fn field(input: &str) -> IResult<&str, RawEventLine<'_>> { + terminated( + ( + take_while1(is_name_char), + opt(preceded( + take_while_m_n(1, 1, is_colon), + preceded(opt(take_while_m_n(1, 1, is_space)), take_while(is_any_char)), + )), + ), + end_of_line, + ) + .parse(input) + .map(|(input, (field, data))| (input, RawEventLine::Field(field, data))) +} + +#[inline] +fn empty(input: &str) -> IResult<&str, RawEventLine<'_>> { + end_of_line(input).map(|(i, _)| (i, RawEventLine::Empty)) +} + +pub fn line(input: &str) -> IResult<&str, RawEventLine<'_>> { + alt((comment, field, empty)).parse(input) +} diff --git a/crates/forge_eventsource_stream/src/traits.rs b/crates/forge_eventsource_stream/src/traits.rs new file mode 100644 index 0000000000000000000000000000000000000000..b5c08341e97bc9f60f4cf61380bdc2ae96fea73e --- /dev/null +++ b/crates/forge_eventsource_stream/src/traits.rs @@ -0,0 +1,19 @@ +use futures_core::stream::Stream; + +use crate::event_stream::EventStream; + +/// Main entrypoint for creating [`crate::Event`] streams +pub trait Eventsource: Sized { + /// Create an event stream from a stream of bytes + fn eventsource(self) -> EventStream; +} + +impl Eventsource for S +where + S: Stream>, + B: AsRef<[u8]>, +{ + fn eventsource(self) -> EventStream { + EventStream::new(self) + } +} diff --git a/crates/forge_eventsource_stream/src/utf8_stream.rs b/crates/forge_eventsource_stream/src/utf8_stream.rs new file mode 100644 index 0000000000000000000000000000000000000000..8d6bfb4b1fbe1f2a41cce71012aef725c3f4cb00 --- /dev/null +++ b/crates/forge_eventsource_stream/src/utf8_stream.rs @@ -0,0 +1,175 @@ +#[cfg(not(feature = "std"))] +use alloc::{ + string::{FromUtf8Error, String}, + vec::Vec, +}; +use core::pin::Pin; +#[cfg(feature = "std")] +use std::string::FromUtf8Error; + +use futures_core::stream::Stream; +use futures_core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { +pub struct Utf8Stream { + #[pin] + stream: S, + buffer: Vec, + terminated: bool, +} +} + +impl Utf8Stream { + pub fn new(stream: S) -> Self { + Self { stream, buffer: Vec::new(), terminated: false } + } +} + +#[derive(Debug, PartialEq)] +pub enum Utf8StreamError { + Utf8(FromUtf8Error), + Transport(E), +} + +impl From for Utf8StreamError { + fn from(err: FromUtf8Error) -> Self { + Self::Utf8(err) + } +} + +impl Stream for Utf8Stream +where + S: Stream>, + B: AsRef<[u8]>, +{ + type Item = Result>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll> { + let this = self.project(); + if *this.terminated { + return Poll::Ready(None); + } + match this.stream.poll_next(cx) { + Poll::Ready(Some(Ok(bytes))) => { + this.buffer.extend_from_slice(bytes.as_ref()); + let bytes = core::mem::take(this.buffer); + match String::from_utf8(bytes) { + Ok(string) => Poll::Ready(Some(Ok(string))), + Err(err) => { + let valid_size = err.utf8_error().valid_up_to(); + let mut bytes = err.into_bytes(); + let rem = bytes.split_off(valid_size); + *this.buffer = rem; + Poll::Ready(Some(Ok(unsafe { String::from_utf8_unchecked(bytes) }))) + } + } + } + Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(Utf8StreamError::Transport(err)))), + Poll::Ready(None) => { + *this.terminated = true; + if this.buffer.is_empty() { + Poll::Ready(None) + } else { + Poll::Ready(Some( + String::from_utf8(core::mem::take(this.buffer)) + .map_err(Utf8StreamError::Utf8), + )) + } + } + Poll::Pending => Poll::Pending, + } + } +} + +#[cfg(test)] +mod tests { + use futures::prelude::*; + + use super::*; + + #[tokio::test] + async fn valid_streams() { + assert_eq!( + Utf8Stream::new(futures::stream::iter(vec![Ok::<_, ()>(b"Hello, world!")])) + .try_collect::>() + .await + .unwrap(), + vec!["Hello, world!"] + ); + assert_eq!( + Utf8Stream::new(futures::stream::iter(vec![Ok::<_, ()>("Hello, world!")])) + .try_collect::>() + .await + .unwrap(), + vec!["Hello, world!"] + ); + assert_eq!( + Utf8Stream::new(futures::stream::iter(vec![Ok::<_, ()>("")])) + .try_collect::>() + .await + .unwrap(), + vec![""] + ); + assert_eq!( + Utf8Stream::new(futures::stream::iter(vec![ + Ok::<_, ()>("Hello"), + Ok::<_, ()>(", world!") + ])) + .try_collect::>() + .await + .unwrap(), + vec!["Hello", ", world!"] + ); + assert_eq!( + Utf8Stream::new(futures::stream::iter(vec![Ok::<_, ()>(vec![ + 240, 159, 145, 141 + ]),])) + .try_collect::>() + .await + .unwrap(), + vec!["👍"] + ); + assert_eq!( + Utf8Stream::new(futures::stream::iter(vec![ + Ok::<_, ()>(vec![240, 159]), + Ok::<_, ()>(vec![145, 141]) + ])) + .try_collect::>() + .await + .unwrap(), + vec!["", "👍"] + ); + assert_eq!( + Utf8Stream::new(futures::stream::iter(vec![ + Ok::<_, ()>(vec![240, 159]), + Ok::<_, ()>(vec![145, 141, 240, 159, 145, 141]) + ])) + .try_collect::>() + .await + .unwrap(), + vec!["", "👍👍"] + ); + } + + #[tokio::test] + async fn invalid_streams() { + let results = Utf8Stream::new(futures::stream::iter(vec![Ok::<_, ()>(vec![240, 159])])) + .collect::>() + .await; + assert_eq!(results.len(), 2); + assert_eq!(results[0], Ok("".to_string())); + assert!(matches!(results[1], Err(Utf8StreamError::Utf8(_)))); + + let results = Utf8Stream::new(futures::stream::iter(vec![ + Ok::<_, ()>(vec![240, 159]), + Ok::<_, ()>(vec![145, 141, 240, 159, 145]), + ])) + .collect::>() + .await; + assert_eq!(results.len(), 3); + assert_eq!(results[0], Ok("".to_string())); + assert_eq!(results[1], Ok("👍".to_string())); + assert!(matches!(results[2], Err(Utf8StreamError::Utf8(_)))); + } +} diff --git a/crates/forge_fs/src/binary_detection.rs b/crates/forge_fs/src/binary_detection.rs new file mode 100644 index 0000000000000000000000000000000000000000..d2c15213c7db4d661ef9ddea179b9e881cfffa30 --- /dev/null +++ b/crates/forge_fs/src/binary_detection.rs @@ -0,0 +1,160 @@ +use anyhow::Result; +use tokio::fs::File; +use tokio::io::AsyncReadExt; + +/// Detected encoding types based on BOM analysis +#[derive(Debug, Clone, PartialEq, Eq)] +enum Encoding { + Utf8WithBom, + Utf16BE, + Utf16LE, +} + +impl Encoding { + /// Detects BOM (Byte Order Mark) patterns + pub fn detect(buffer: &[u8], bytes_read: usize) -> Option { + match (buffer.first(), buffer.get(1), buffer.get(2)) { + (Some(&0xEF), Some(&0xBB), Some(&0xBF)) if bytes_read >= 3 => { + Some(Encoding::Utf8WithBom) + } + (Some(&0xFE), Some(&0xFF), _) if bytes_read >= 2 => Some(Encoding::Utf16BE), + (Some(&0xFF), Some(&0xFE), _) if bytes_read >= 2 => Some(Encoding::Utf16LE), + _ => None, + } + } +} + +/// Detects if a file is binary by analyzing its content +pub async fn is_binary>(path: P) -> Result { + use anyhow::Context; + let path_ref = path.as_ref(); + let mut file = File::open(path_ref) + .await + .with_context(|| format!("Failed to open file {}", path_ref.display()))?; + + let mut buffer = vec![0u8; 512]; + let bytes_read = file.read(&mut buffer).await?; + buffer.truncate(bytes_read); + + Ok(is_binary_internal(&buffer, bytes_read)) +} + +/// Detects encoding and binary status from a buffer +fn is_binary_internal(buffer: &[u8], bytes_read: usize) -> bool { + // Always first check for BOM to find out about encoding + let encoding = Encoding::detect(buffer, bytes_read); + + // Detect 0 bytes to see if file is binary or UTF-16 LE/BE + // unless we already know that this file has a UTF-16 encoding + let mut seems_binary = false; + if encoding != Some(Encoding::Utf16BE) + && encoding != Some(Encoding::Utf16LE) + && !buffer.is_empty() + { + let mut could_be_utf16le = true; // e.g. 0xAA 0x00 + let mut could_be_utf16be = true; // e.g. 0x00 0xAA + let mut contains_zero_byte = false; + + // This is a simplified guess to detect UTF-16 BE or LE by just checking if + // the first 512 bytes have the 0-byte at a specific location. For UTF-16 LE + // this would be the odd byte index and for UTF-16 BE the even one. + // Note: this can produce false positives (a binary file that uses a 2-byte + // encoding of the same format as UTF-16) and false negatives (a UTF-16 file + // that is using 4 bytes to encode a character). + const ZERO_BYTE_DETECTION_BUFFER_MAX_LEN: usize = 512; + for (i, &byte) in buffer + .iter() + .enumerate() + .take(bytes_read.min(ZERO_BYTE_DETECTION_BUFFER_MAX_LEN)) + { + let is_endian = i % 2 == 1; // assume 2-byte sequences typical for UTF-16 + let is_zero_byte = byte == 0; + + if is_zero_byte { + contains_zero_byte = true; + } + + // UTF-16 LE: expect e.g. 0xAA 0x00 + if could_be_utf16le && (is_endian && !is_zero_byte || !is_endian && is_zero_byte) { + could_be_utf16le = false; + } + + // UTF-16 BE: expect e.g. 0x00 0xAA + if could_be_utf16be && (is_endian && is_zero_byte || !is_endian && !is_zero_byte) { + could_be_utf16be = false; + } + + // Return if this is neither UTF16-LE nor UTF16-BE and thus treat as binary + if is_zero_byte && !could_be_utf16le && !could_be_utf16be { + break; + } + } + + // Handle case of 0-byte included + if contains_zero_byte && !could_be_utf16le && !could_be_utf16be { + seems_binary = true; + } + } + + seems_binary +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use tempfile::NamedTempFile; + use tokio::fs; + + use super::*; + + async fn create_test_file_fixture(content: &[u8]) -> Result { + let file = NamedTempFile::new()?; + fs::write(file.path(), content).await?; + Ok(file) + } + + #[tokio::test] + async fn test_empty_file_is_binary() -> Result<()> { + let fixture = create_test_file_fixture(&[]).await?; + let actual = is_binary(fixture.path()).await?; + let expected = false; + assert_eq!(actual, expected); + Ok(()) + } + + #[tokio::test] + async fn test_plain_text_file_is_binary() -> Result<()> { + let fixture = create_test_file_fixture(b"Hello, world!").await?; + let actual = is_binary(fixture.path()).await?; + let expected = false; + assert_eq!(actual, expected); + Ok(()) + } + + #[tokio::test] + async fn test_binary_file_with_zero_bytes() -> Result<()> { + let content = vec![ + 0x48, 0x65, 0x6C, 0x6C, 0x6F, 0x00, 0x57, 0x6F, 0x72, 0x6C, 0x64, + ]; + let fixture = create_test_file_fixture(&content).await?; + let actual = is_binary(fixture.path()).await?; + let expected = true; + assert_eq!(actual, expected); + Ok(()) + } + + #[tokio::test] + async fn test_buffer_limit_512_bytes() -> Result<()> { + // Create content larger than 512 bytes with zero byte at position 600 + let mut content = vec![0x48; 600]; // 'H' repeated 600 times + content[599] = 0x00; // Zero byte beyond 512 byte limit + + let fixture = create_test_file_fixture(&content).await?; + let actual = is_binary(fixture.path()).await?; + + // Should not detect as binary because zero byte is beyond 512-byte limit + let expected = false; + assert_eq!(actual, expected); + Ok(()) + } +} diff --git a/crates/forge_fs/src/error.rs b/crates/forge_fs/src/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..5f2291da35aae540628e4800d25fb8ece4bd0484 --- /dev/null +++ b/crates/forge_fs/src/error.rs @@ -0,0 +1,25 @@ +use std::string::FromUtf8Error; + +use thiserror::Error; + +/// Error type for file operations +#[derive(Error, Debug)] +pub enum Error { + #[error("Binary files are not supported. File detected as {0}")] + BinaryFileNotSupported(String), + + #[error("Start position {start} is beyond the file size of {total} characters")] + StartBeyondFileSize { start: u64, total: u64 }, + + #[error("Start position {start} and end position {end} must be 1-based (inclusive)")] + IndexStartingWithZero { start: u64, end: u64 }, + + #[error("Start position {start} is greater than end position {end}")] + StartGreaterThanEnd { start: u64, end: u64 }, + + #[error("UTF-8 validation failed: {0}")] + Utf8ValidationFailed(#[from] FromUtf8Error), + + #[error("I/O error: {0}")] + IoError(#[from] std::io::Error), +} diff --git a/crates/forge_fs/src/file_size.rs b/crates/forge_fs/src/file_size.rs new file mode 100644 index 0000000000000000000000000000000000000000..5c7c4866b916f66c6a26cc6c45cd7891a1912743 --- /dev/null +++ b/crates/forge_fs/src/file_size.rs @@ -0,0 +1,17 @@ +use std::path::Path; + +use anyhow::{Context, Result}; + +impl crate::ForgeFS { + /// Gets file size without reading the entire file + pub async fn file_size>(path: T) -> Result { + let metadata = tokio::fs::metadata(path.as_ref()).await.with_context(|| { + format!( + "Failed to get metadata for file {}", + path.as_ref().display() + ) + })?; + + Ok(metadata.len()) + } +} diff --git a/crates/forge_fs/src/is_binary.rs b/crates/forge_fs/src/is_binary.rs new file mode 100644 index 0000000000000000000000000000000000000000..6dfe303cf5a3950851bdb668438a6be79179355c --- /dev/null +++ b/crates/forge_fs/src/is_binary.rs @@ -0,0 +1,104 @@ +use anyhow::Result; +use tokio::fs::File; +use tokio::io::AsyncReadExt; + +impl crate::ForgeFS { + /// Checks if a file is binary by examining its content. + /// This version takes a path and opens the file itself. + #[cfg(test)] + async fn is_binary_path>(path: T) -> Result<(bool, String)> { + use anyhow::Context; + + let path_ref = path.as_ref(); + let mut file = File::open(path_ref) + .await + .with_context(|| format!("Failed to open file {}", path_ref.display()))?; + + Self::is_binary(&mut file).await + } + + /// Checks if a file is binary by examining its content. + /// This version takes an already opened file handle, allowing for reuse + /// of the same file handle across multiple operations. + /// This is a crate-private implementation detail. + pub(crate) async fn is_binary(file: &mut File) -> Result<(bool, String)> { + // Read sample data + let mut sample = vec![0; 8192]; + let bytes_read = file.read(&mut sample).await?; + sample.truncate(bytes_read); + + // Handle empty files + if bytes_read == 0 { + return Ok((true, "Empty file".into())); + } + + // Get file type info + let is_text = match infer::get(&sample) { + Some(info) => matches!( + info.matcher_type(), + infer::MatcherType::Text | infer::MatcherType::Doc + ), + None => true, // Assume text if type can't be determined + }; + + let description = infer::get(&sample) + .map(|info| info.mime_type().to_string()) + .unwrap_or_else(|| "Text file (no specific format detected)".into()); + + Ok((is_text, description)) + } +} + +#[cfg(test)] +mod test { + use anyhow::Result; + use tokio::fs; + + async fn create_test_file(content: &[u8]) -> Result { + let file = tempfile::NamedTempFile::new()?; + fs::write(file.path(), content).await?; + Ok(file) + } + + #[tokio::test] + async fn test_is_binary_file() -> Result<()> { + // Test text file + let text_file = create_test_file(b"Hello, world!").await?; + let (is_text_or_doc, _) = crate::ForgeFS::is_binary_path(text_file.path()).await?; + assert!(is_text_or_doc, "Text file should be identified as text"); + + // Test binary data + let binary_content = vec![0, 1, 2, 3, 0, 0, 0, 0, 5, 6, 7, 8]; + let binary_file = create_test_file(&binary_content).await?; + let (is_text_or_doc, file_type) = + crate::ForgeFS::is_binary_path(binary_file.path()).await?; + + if !is_text_or_doc { + assert!( + file_type.contains("binary") || !file_type.contains("text"), + "Binary file type description should indicate binary" + ); + } + + // Test PNG file + let png_header = [ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, + 0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, 0x00, 0x00, + 0x00, + ]; + let png_file = create_test_file(&png_header).await?; + let (is_text_or_doc, file_type) = crate::ForgeFS::is_binary_path(png_file.path()).await?; + assert!(!is_text_or_doc, "PNG file should be identified as binary"); + assert!( + file_type.contains("image/png"), + "PNG file type should be correctly identified" + ); + + // Test empty file + let empty_file = create_test_file(&[]).await?; + let (is_text_or_doc, _) = crate::ForgeFS::is_binary_path(empty_file.path()).await?; + assert!(is_text_or_doc, "Empty file should be considered text"); + + Ok(()) + } +} diff --git a/crates/forge_fs/src/lib.rs b/crates/forge_fs/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..e1d570ee9d9ebd16dc1850b40e6a586c7c7ddd43 --- /dev/null +++ b/crates/forge_fs/src/lib.rs @@ -0,0 +1,36 @@ +//! # ForgeFS +//! +//! A file system abstraction layer that standardizes error handling for file +//! operations. +//! +//! ForgeFS wraps tokio's filesystem operations with consistent error context +//! using anyhow::Context. Each method provides standardized error messages in +//! the format "Failed to [operation] [path]", ensuring uniform error reporting +//! throughout the application while preserving the original error cause. + +mod binary_detection; +mod error; +mod file_size; +mod is_binary; +mod meta; +mod read; +mod read_range; +mod write; + +pub use crate::binary_detection::is_binary; +pub use crate::error::Error; + +/// ForgeFS provides a standardized interface for file system operations +/// with consistent error handling. +#[derive(Debug)] +pub struct ForgeFS; + +impl ForgeFS { + /// Computes a SHA-256 hash of the given string content. + pub(crate) fn compute_hash(content: &str) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(content.as_bytes()); + hex::encode(hasher.finalize()) + } +} diff --git a/crates/forge_fs/src/meta.rs b/crates/forge_fs/src/meta.rs new file mode 100644 index 0000000000000000000000000000000000000000..bd5348d832b8fe56320c71ea47e5324050505901 --- /dev/null +++ b/crates/forge_fs/src/meta.rs @@ -0,0 +1,25 @@ +use std::path::Path; + +use anyhow::{Context, Result}; + +use crate::is_binary; + +impl crate::ForgeFS { + pub fn exists>(path: T) -> bool { + path.as_ref().exists() + } + + pub async fn is_binary_file>(path: T) -> anyhow::Result { + is_binary(path).await + } + + pub fn is_file>(path: T) -> bool { + path.as_ref().is_file() + } + + pub async fn read_dir>(path: T) -> Result { + tokio::fs::read_dir(path.as_ref()) + .await + .with_context(|| format!("Failed to read directory {}", path.as_ref().display())) + } +} diff --git a/crates/forge_fs/src/read.rs b/crates/forge_fs/src/read.rs new file mode 100644 index 0000000000000000000000000000000000000000..b96c3c82e42fccda84d928ad48b15d550623fa90 --- /dev/null +++ b/crates/forge_fs/src/read.rs @@ -0,0 +1,24 @@ +use std::path::Path; + +use anyhow::{Context, Result}; +use bstr::ByteSlice; + +impl crate::ForgeFS { + pub async fn read_utf8>(path: T) -> Result { + Self::read(path) + .await + .map(|bytes| bytes.to_str_lossy().to_string()) + } + + pub async fn read>(path: T) -> Result> { + tokio::fs::read(path.as_ref()) + .await + .with_context(|| format!("Failed to read file {}", path.as_ref().display())) + } + + pub async fn read_to_string>(path: T) -> Result { + tokio::fs::read_to_string(path.as_ref()) + .await + .with_context(|| format!("Failed to read file as string {}", path.as_ref().display())) + } +} diff --git a/crates/forge_fs/src/read_range.rs b/crates/forge_fs/src/read_range.rs new file mode 100644 index 0000000000000000000000000000000000000000..ed48f5de1bf8632355c4d5cd2f9b81b8f7a90401 --- /dev/null +++ b/crates/forge_fs/src/read_range.rs @@ -0,0 +1,335 @@ +use std::cmp; +use std::path::Path; + +use anyhow::{Context, Result}; +use bstr::ByteSlice; +use forge_domain::FileInfo; + +use crate::error::Error; + +impl crate::ForgeFS { + /// Reads a specific range of lines from a file. + /// + /// # Arguments + /// * `path` - Path to the file to read + /// * `start_line` - Starting line number (1-based, inclusive) + /// * `end_line` - Ending line number (1-based, inclusive) + /// + /// Returns a tuple containing: + /// - The range content as a UTF-8 string. + /// - FileInfo containing metadata about the read operation including line + /// positions and a `content_hash` computed from the **full** file + /// content, so callers can store a hash that matches what a subsequent + /// whole-file read would produce (used by the external-change detector). + pub async fn read_range_utf8>( + path: T, + start_line: u64, + end_line: u64, + ) -> Result<(String, FileInfo)> { + let path_ref = path.as_ref(); + + // Basic validation + if start_line > end_line { + return Err(Error::StartGreaterThanEnd { start: start_line, end: end_line }.into()); + } + + // Open and check if file is binary + let mut file = tokio::fs::File::open(path_ref) + .await + .with_context(|| format!("Failed to open file {}", path_ref.display()))?; + + if start_line == 0 || end_line == 0 { + return Err(Error::IndexStartingWithZero { start: start_line, end: end_line }.into()); + } + + let (is_text, file_type) = Self::is_binary(&mut file).await?; + if !is_text { + return Err(Error::BinaryFileNotSupported(file_type).into()); + } + + // Read file content + let content = tokio::fs::read(path_ref) + .await + .with_context(|| format!("Failed to read file content from {}", path_ref.display()))?; + let content = content.to_str_lossy(); + + // Hash the full file content so callers get a stable, whole-file hash + // that matches what the external-change detector reads back from disk. + let content_hash = crate::ForgeFS::compute_hash(content.as_ref()); + + if start_line < 2 && content.is_empty() { + // If the file is empty, return empty content + return Ok(( + String::new(), + FileInfo::new(start_line, end_line, 0, content_hash), + )); + } + // Split into lines + let lines: Vec<&str> = content.lines().collect(); + let total_lines = lines.len() as u64; + + // Convert to 0-based indexing + let start_pos = start_line.saturating_sub(1); + let mut end_pos = end_line.saturating_sub(1); + + // Validate start position + if start_pos >= total_lines { + return Err( + Error::StartBeyondFileSize { start: start_line, total: total_lines }.into(), + ); + } + + // Cap end position at last line + end_pos = cmp::min(end_pos, total_lines - 1); + + // Calculate actual end line (1-based) that was used + let actual_end_line = cmp::min(end_line, total_lines); + + let info = FileInfo::new(start_line, actual_end_line, total_lines, content_hash); + + // Extract requested lines + let result_content = if start_pos == 0 && end_pos == total_lines - 1 { + content.to_string() // Return full content if requesting entire file + } else { + lines + .get(start_pos as usize..=end_pos as usize) + .map(|slice| slice.join("\n")) + .unwrap_or_default() + }; + + Ok((result_content, info)) + } +} + +#[cfg(test)] +mod test { + use anyhow::Result; + use pretty_assertions::assert_eq; + use tokio::fs; + + // Helper to create a temporary file with test content + async fn create_test_file(content: &str) -> Result { + let file = tempfile::NamedTempFile::new()?; + fs::write(file.path(), content).await?; + Ok(file) + } + + #[tokio::test] + async fn test_read_range_utf8() -> Result<()> { + let content = + "Line 1\nLine 2\nLine 3\nLine 4\nLine 5\nLine 6\nLine 7\nLine 8\nLine 9\nLine 10"; + let file = create_test_file(content).await?; + let full_hash = crate::ForgeFS::compute_hash(content); + + // Test reading a range of lines + let (result, info) = crate::ForgeFS::read_range_utf8(file.path(), 2, 5).await?; + assert_eq!(result, "Line 2\nLine 3\nLine 4\nLine 5"); + assert_eq!(info.start_line, 2); + assert_eq!(info.end_line, 5); + assert_eq!(info.total_lines, 10); + assert_eq!(info.content_hash, full_hash); + + // Test reading from start + let (result, info) = crate::ForgeFS::read_range_utf8(file.path(), 1, 3).await?; + assert_eq!(result, "Line 1\nLine 2\nLine 3"); + assert_eq!(info.start_line, 1); + assert_eq!(info.end_line, 3); + assert_eq!(info.content_hash, full_hash); + + // Test reading to end + let (result, info) = crate::ForgeFS::read_range_utf8(file.path(), 8, 10).await?; + assert_eq!(result, "Line 8\nLine 9\nLine 10"); + assert_eq!(info.start_line, 8); + assert_eq!(info.end_line, 10); + assert_eq!(info.content_hash, full_hash); + + // Test reading entire file + let (result, info) = crate::ForgeFS::read_range_utf8(file.path(), 1, 10).await?; + assert_eq!(result, content); + assert_eq!(info.start_line, 1); + assert_eq!(info.end_line, 10); + assert_eq!(info.content_hash, full_hash); + + // Test single line + let (result, info) = crate::ForgeFS::read_range_utf8(file.path(), 5, 5).await?; + assert_eq!(result, "Line 5"); + assert_eq!(info.start_line, 5); + assert_eq!(info.end_line, 5); + assert_eq!(info.content_hash, full_hash); + + // Test first line specifically + let (result, info) = crate::ForgeFS::read_range_utf8(file.path(), 1, 1).await?; + assert_eq!(result, "Line 1"); + assert_eq!(info.start_line, 1); + assert_eq!(info.end_line, 1); + assert_eq!(info.total_lines, 10); + assert_eq!(info.content_hash, full_hash); + + // Test invalid ranges + assert!( + crate::ForgeFS::read_range_utf8(file.path(), 8, 5) + .await + .is_err() + ); + assert!( + crate::ForgeFS::read_range_utf8(file.path(), 15, 10) + .await + .is_err() + ); + assert!( + crate::ForgeFS::read_range_utf8(file.path(), 0, 5) + .await + .is_err() + ); + + Ok(()) + } + + #[tokio::test] + async fn test_utf8_multi_line_handling() -> Result<()> { + let content = "Hello world!\nこんにちは 世界!\nПривет мир!\nBonjour le monde!"; + let file = create_test_file(content).await?; + + // Test reading a range that includes multi-byte characters + let (result, info) = crate::ForgeFS::read_range_utf8(file.path(), 2, 3).await?; + assert_eq!(result, "こんにちは 世界!\nПривет мир!"); + assert_eq!(info.start_line, 2); + assert_eq!(info.end_line, 3); + + Ok(()) + } + + #[tokio::test] + async fn test_invalid_utf8_handling() -> Result<()> { + let content = b"Hello world!\nValid line\n\xFF\xFE\xFD Invalid UTF-8\nAnother valid line"; + let file = tempfile::NamedTempFile::new()?; + fs::write(file.path(), content).await?; + + // Attempt to read the file shouldn't fail with invalid UTF-8 error + let result = crate::ForgeFS::read_range_utf8(&file.path(), 1, 4).await; + assert!(result.is_ok()); + Ok(()) + } + + #[tokio::test] + async fn test_end_line_capping() -> Result<()> { + let content = "Line 1\nLine 2\nLine 3\nLine 4\nLine 5"; + let file = create_test_file(content).await?; + + // Test: end_line = total_lines (exact match) + let (actual_content, actual_info) = + crate::ForgeFS::read_range_utf8(file.path(), 1, 5).await?; + assert_eq!(actual_content, content); + assert_eq!( + ( + actual_info.start_line, + actual_info.end_line, + actual_info.total_lines + ), + (1, 5, 5) + ); + + // Test: end_line = total_lines + 1 (one beyond) + let (actual_content, actual_info) = + crate::ForgeFS::read_range_utf8(file.path(), 1, 6).await?; + assert_eq!(actual_content, content); + assert_eq!( + ( + actual_info.start_line, + actual_info.end_line, + actual_info.total_lines + ), + (1, 5, 5) + ); + + // Test: end_line >> total_lines (far beyond) + let (actual_content, actual_info) = + crate::ForgeFS::read_range_utf8(file.path(), 1, 10000).await?; + assert_eq!(actual_content, content); + assert_eq!( + ( + actual_info.start_line, + actual_info.end_line, + actual_info.total_lines + ), + (1, 5, 5) + ); + + // Test: range starting in the middle with excessive end_line + let (actual_content, actual_info) = + crate::ForgeFS::read_range_utf8(file.path(), 3, 100).await?; + assert_eq!(actual_content, "Line 3\nLine 4\nLine 5"); + assert_eq!( + ( + actual_info.start_line, + actual_info.end_line, + actual_info.total_lines + ), + (3, 5, 5) + ); + + // Test: reading last line with excessive end_line + let (actual_content, actual_info) = + crate::ForgeFS::read_range_utf8(file.path(), 5, 100).await?; + assert_eq!(actual_content, "Line 5"); + assert_eq!( + ( + actual_info.start_line, + actual_info.end_line, + actual_info.total_lines + ), + (5, 5, 5) + ); + + Ok(()) + } + + #[tokio::test] + async fn test_large_file_ranges() -> Result<()> { + // Create a 5000-line file to test various range scenarios + let lines: Vec = (1..=5000).map(|i| format!("Line {}", i)).collect(); + let content = lines.join("\n"); + let file = create_test_file(&content).await?; + + // Test: range within file bounds + let (actual_content, actual_info) = + crate::ForgeFS::read_range_utf8(file.path(), 1, 50).await?; + assert_eq!(actual_content.lines().count(), 50); + assert_eq!( + ( + actual_info.start_line, + actual_info.end_line, + actual_info.total_lines + ), + (1, 50, 5000) + ); + + // Test: range in the middle + let (actual_content, actual_info) = + crate::ForgeFS::read_range_utf8(file.path(), 2500, 3500).await?; + assert_eq!(actual_content.lines().count(), 1001); // Lines 2500-3500 inclusive + assert_eq!( + ( + actual_info.start_line, + actual_info.end_line, + actual_info.total_lines + ), + (2500, 3500, 5000) + ); + + // Test: end beyond file bounds (simulates max_read_size limiting) + let (actual_content, actual_info) = + crate::ForgeFS::read_range_utf8(file.path(), 4990, 6000).await?; + assert_eq!(actual_content.lines().count(), 11); // Lines 4990-5000 + assert_eq!( + ( + actual_info.start_line, + actual_info.end_line, + actual_info.total_lines + ), + (4990, 5000, 5000) + ); + + Ok(()) + } +} diff --git a/crates/forge_fs/src/write.rs b/crates/forge_fs/src/write.rs new file mode 100644 index 0000000000000000000000000000000000000000..7972e6e94c6cf40cbab5c0d895c38efcc1c25ef9 --- /dev/null +++ b/crates/forge_fs/src/write.rs @@ -0,0 +1,42 @@ +use std::path::Path; + +use anyhow::{Context, Result}; +use tokio::io::AsyncWriteExt as _; + +impl crate::ForgeFS { + pub async fn create_dir_all>(path: T) -> Result<()> { + tokio::fs::create_dir_all(path.as_ref()) + .await + .with_context(|| format!("Failed to create dir {}", path.as_ref().display())) + } + + pub async fn write, U: AsRef<[u8]>>(path: T, contents: U) -> Result<()> { + tokio::fs::write(path.as_ref(), contents) + .await + .with_context(|| format!("Failed to write file {}", path.as_ref().display())) + } + + /// Appends content to an existing file, or creates it if it does not exist. + pub async fn append, U: AsRef<[u8]>>(path: T, contents: U) -> Result<()> { + let mut file = tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path.as_ref()) + .await + .with_context(|| { + format!( + "Failed to open file for appending {}", + path.as_ref().display() + ) + })?; + file.write_all(contents.as_ref()) + .await + .with_context(|| format!("Failed to append to file {}", path.as_ref().display())) + } + + pub async fn remove_file>(path: T) -> Result<()> { + tokio::fs::remove_file(path.as_ref()) + .await + .with_context(|| format!("Failed to remove file {}", path.as_ref().display())) + } +} diff --git a/crates/forge_main/src/banner b/crates/forge_main/src/banner new file mode 100644 index 0000000000000000000000000000000000000000..51438f09ca2bd4b37a11ecd1bc8af6218278a67c --- /dev/null +++ b/crates/forge_main/src/banner @@ -0,0 +1,6 @@ + _____ +| ___|__ _ __ __ _ ___ +| |_ / _ \| '__/ _` |/ _ \ +| _| (_) | | | (_| | __/ +|_| \___/|_| \__, |\___| + |___/ diff --git a/crates/forge_main/src/banner.rs b/crates/forge_main/src/banner.rs new file mode 100644 index 0000000000000000000000000000000000000000..1827ddaa529ad0f23a572711a1c16956017e97a6 --- /dev/null +++ b/crates/forge_main/src/banner.rs @@ -0,0 +1,138 @@ +use std::{fmt, io}; + +use colored::Colorize; +use forge_tracker::VERSION; + +const BANNER: &str = include_str!("banner"); + +/// Renders messages into a styled box with border characters. +struct DisplayBox { + messages: Vec, +} + +impl DisplayBox { + /// Creates a new Box with the given messages. + fn new(messages: Vec) -> Self { + Self { messages } + } +} + +impl fmt::Display for DisplayBox { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let visible_len = |s: &str| console::measure_text_width(s); + let width: usize = self + .messages + .iter() + .map(|s| visible_len(s)) + .max() + .unwrap_or(0) + + 4; + let top = format!("┌{}┐", "─".repeat(width.saturating_sub(2))); + let bottom = format!("└{}┘", "─".repeat(width.saturating_sub(2))); + let fmt_line = |s: &str| { + let padding = width.saturating_sub(4).saturating_sub(visible_len(s)); + format!("│ {}{} │", s, " ".repeat(padding)) + }; + + writeln!(f, "{}", top)?; + for msg in &self.messages { + writeln!(f, "{}", fmt_line(msg))?; + } + write!(f, "{}", bottom) + } +} + +/// Displays the banner with version and command tips. +/// +/// # Arguments +/// +/// * `cli_mode` - If true, shows CLI-relevant commands. Both interactive and +/// CLI modes use `:` as the canonical command prefix. +/// +/// # Environment Variables +/// +/// * `FORGE_BANNER` - Optional custom banner text to display instead of the +/// default +pub fn display(cli_mode: bool) -> io::Result<()> { + // Check for custom banner via environment variable + let mut banner = std::env::var("FORGE_BANNER") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| BANNER.to_string()); + + // Always show version + let version_label = ("Version:", VERSION); + + // Build tips based on mode + let tips: Vec<(&str, &str)> = if cli_mode { + // CLI mode: only show relevant commands + vec![ + ("New conversation:", ":new"), + ("Get started:", ":info, :conversation"), + ("Switch model:", ":model"), + ("Switch provider:", ":provider"), + ("Switch agent:", ": e.g. :forge or :muse"), + ] + } else { + // Interactive mode: show all commands + vec![ + ("New conversation:", ":new"), + ("Get started:", ":info, :usage, :help, :conversation"), + ("Switch model:", ":model"), + ("Switch agent:", ":forge or :muse or :agent"), + ("Update:", ":update"), + ("Quit:", ":exit or "), + ] + }; + + // Build labels array with version and tips + let labels: Vec<(&str, &str)> = std::iter::once(version_label).chain(tips).collect(); + + // Calculate the width of the longest label key for alignment + let max_width = labels.iter().map(|(key, _)| key.len()).max().unwrap_or(0); + + // Add all lines with right-aligned label keys and their values + for (key, value) in &labels { + banner.push_str( + format!( + "\n{}{}", + format!("{key:>max_width$} ").dimmed(), + value.cyan() + ) + .as_str(), + ); + } + + println!("{banner}\n"); + + // Encourage zsh integration after the banner + if !cli_mode { + display_zsh_encouragement(); + } + + Ok(()) +} + +/// Encourages users to use the zsh plugin for a better experience. +fn display_zsh_encouragement() { + let tip = DisplayBox::new(vec![ + format!( + "{} {}", + "TIP:".bold().yellow(), + "For the best experience, use our zsh plugin!".bold() + ), + format!( + "{} {} {}", + "·".dimmed(), + "Set up forge via our zsh plugin:".dimmed(), + "forge zsh setup".bold().green(), + ), + format!( + "{} {} {}", + "·".dimmed(), + "Learn more:".dimmed(), + "https://forgecode.dev/docs/zsh-support".cyan() + ), + ]); + println!("{}", tip); +} diff --git a/crates/forge_main/src/cli.rs b/crates/forge_main/src/cli.rs new file mode 100644 index 0000000000000000000000000000000000000000..1068f77590ee175371ea2932d2186a07b0e45f7f --- /dev/null +++ b/crates/forge_main/src/cli.rs @@ -0,0 +1,2037 @@ +//! NOTE: Always use singular names for commands and subcommands. +//! For example: `forge provider login` instead of `forge providers login`. +//! +//! NOTE: With every change to this CLI structure, verify that the ZSH plugin +//! remains compatible. The plugin at `shell-plugin/forge.plugin.zsh` implements +//! shell completion and command shortcuts that depend on the CLI structure. + +use std::path::PathBuf; + +use clap::{Parser, Subcommand, ValueEnum}; +use forge_domain::{AgentId, ConversationId, Effort, ModelId, ProviderId}; + +#[derive(Parser)] +#[command(version = env!("CARGO_PKG_VERSION"))] +pub struct Cli { + /// Direct prompt to process without entering interactive mode. + /// + /// When provided, executes a single command and exits instead of starting + /// an interactive session. Content can also be piped: `cat prompt.txt | + /// forge`. + #[arg(long, short = 'p', allow_hyphen_values = true)] + pub prompt: Option, + + /// Piped input from stdin (populated internally) + /// + /// This field is automatically populated when content is piped to forge + /// via stdin. It's kept separate from the prompt to allow proper handling + /// as a droppable message. + #[arg(skip)] + pub piped_input: Option, + + /// Path to a JSON file containing the conversation to execute. + #[arg(long)] + pub conversation: Option, + + /// Conversation ID to use for this session. + /// + /// When provided, resumes or continues an existing conversation instead of + /// generating a new conversation ID. + #[arg(long, alias = "cid")] + pub conversation_id: Option, + + /// Working directory to use before starting the session. + /// + /// When provided, changes to this directory before starting forge. + #[arg(long, short = 'C')] + pub directory: Option, + + /// Name for an isolated git worktree to create for experimentation. + #[arg(long)] + pub sandbox: Option, + + /// Enable verbose logging output. + #[arg(long, default_value_t = false)] + pub verbose: bool, + + /// Agent ID to use for this session. + #[arg(long, alias = "aid")] + pub agent: Option, + + /// Top-level subcommands. + #[command(subcommand)] + pub subcommands: Option, + + /// Event to dispatch to the workflow in JSON format. + #[arg(long, short = 'e')] + pub event: Option, +} + +impl Cli { + /// Determines whether the CLI should start in interactive mode. + /// + /// Returns true when no prompt, piped input, or subcommand is provided, + /// indicating the user wants to enter interactive mode. + pub fn is_interactive(&self) -> bool { + self.prompt.is_none() && self.piped_input.is_none() && self.subcommands.is_none() + } +} + +#[derive(Subcommand, Debug, Clone)] +pub enum TopLevelCommand { + /// Manage agents. + Agent(AgentCommandGroup), + + /// Generate shell extension scripts. + #[command(subcommand, alias = "extension")] + Zsh(ZshCommandGroup), + + /// List agents, models, providers, tools, or MCP servers. + List(ListCommandGroup), + + /// Display the banner with version information. + Banner, + + /// Show configuration, active model, and environment status. + Info { + /// Conversation ID for session-specific information. + #[arg(long, alias = "cid")] + conversation_id: Option, + + /// Output in machine-readable format. + #[arg(long)] + porcelain: bool, + }, + + /// Get, set, or list configuration values. + Config(ConfigCommandGroup), + + /// Manage conversation history and state. + #[command(alias = "session")] + Conversation(ConversationCommandGroup), + + /// Generate and optionally commit changes with AI-generated message + Commit(CommitCommandGroup), + + /// Manage Model Context Protocol servers. + Mcp(McpCommandGroup), + + /// Suggest shell commands from natural language. + Suggest { + /// Natural language description of the desired command. + #[arg(allow_hyphen_values = true)] + prompt: String, + }, + + /// Manage API provider authentication. + Provider(ProviderCommandGroup), + + /// Run or list custom commands. + #[command(aliases = ["command", "commands"])] + Cmd(CmdCommandGroup), + + /// Manage workspaces for semantic search. + Workspace(WorkspaceCommandGroup), + + /// Process JSONL data through LLM with schema-constrained tools. + Data(DataCommandGroup), + + /// VS Code integration commands. + #[command(subcommand)] + Vscode(VscodeCommand), + + /// Update forge to the latest version. + Update(UpdateArgs), + + /// Setup zsh integration by updating .zshrc with plugin and theme (alias + /// for `zsh setup`). + Setup, + + /// Run diagnostics on shell environment (alias for `zsh doctor`). + Doctor, + + /// Stream forge log output (defaults to the most recent log file). + Logs(LogsArgs), + + /// Interactive fuzzy item picker. + Select(SelectCommandGroup), +} + +/// Command group for the `forge select` interactive picker. +/// +/// Subcommands provide purpose-built pickers for specific domain types (models, +/// agents, providers, etc.) that fetch data internally and output the selected +/// value. +#[derive(Parser, Debug, Clone)] +pub struct SelectCommandGroup { + #[command(subcommand)] + pub command: SelectCommand, +} + +/// Purpose-built interactive pickers for specific domain types. +/// +/// Each variant fetches data internally through the Forge API, presents an +/// interactive fuzzy picker, and prints the selected value to stdout for the +/// shell plugin to consume. +#[derive(Subcommand, Debug, Clone)] +pub enum SelectCommand { + /// Select a model interactively from all configured providers. + /// + /// Prints the selected model_id on the first line and provider_id on the + /// second line. Prints nothing if the user cancels. + Model { + /// Initial query text pre-filled in the search box. + #[arg(long, short = 'q')] + query: Option, + }, + + /// Select an agent interactively. + /// + /// Prints the selected agent_id on stdout. Prints nothing if the user + /// cancels. + Agent { + /// Initial query text pre-filled in the search box. + #[arg(long, short = 'q')] + query: Option, + }, + + /// Select a provider interactively. + /// + /// Prints the selected provider_id on stdout. Prints nothing if the user + /// cancels. + Provider { + /// Initial query text pre-filled in the search box. + #[arg(long, short = 'q')] + query: Option, + + /// Only show providers that are configured (logged in). + #[arg(long)] + configured: bool, + }, + + /// Select a reasoning effort level interactively. + /// + /// Prints the selected effort level (none, minimal, low, medium, high, + /// xhigh, max) on stdout. Prints nothing if the user cancels. + ReasoningEffort { + /// Initial query text pre-filled in the search box. + #[arg(long, short = 'q')] + query: Option, + }, + + /// Select a command interactively. + /// + /// Prints the selected command name on stdout. Prints nothing if the user + /// cancels. + Command { + /// Initial query text pre-filled in the search box. + #[arg(long, short = 'q')] + query: Option, + }, + + /// Select a conversation interactively with a preview pane. + /// + /// Prints the selected conversation_id on stdout. Prints nothing if the + /// user cancels. + Conversation { + /// Initial query text pre-filled in the search box. + #[arg(long, short = 'q')] + query: Option, + + /// Show child conversations of a parent conversation. + #[arg(long)] + parent: Option, + }, + + /// Select a file interactively with a preview pane. + /// + /// Walks the workspace and presents a fuzzy picker with syntax-highlighted + /// file previews. Prints the selected file path on stdout. Prints nothing + /// if the user cancels. + File { + /// Initial query text pre-filled in the search box. + #[arg(long, short = 'q')] + query: Option, + }, +} + +/// Command group for custom command management. +#[derive(Parser, Debug, Clone)] +pub struct CmdCommandGroup { + #[command(subcommand)] + pub command: CmdCommand, + + /// Conversation ID to execute the command within. + #[arg(long, alias = "cid", global = true)] + pub conversation_id: Option, + + /// Output in machine-readable format. + #[arg(long, global = true)] + pub porcelain: bool, +} + +#[derive(Subcommand, Debug, Clone)] +pub enum CmdCommand { + /// List all available custom commands. + List { + /// Shows only custom commands + #[arg(long)] + custom: bool, + }, + + /// Execute a custom command. + Execute { + /// Name of the custom command to execute, followed by any arguments. + commands: Vec, + }, +} + +/// Command group for agent management. +#[derive(Parser, Debug, Clone)] +pub struct AgentCommandGroup { + #[command(subcommand)] + pub command: AgentCommand, + + /// Output in machine-readable format. + #[arg(long, global = true)] + pub porcelain: bool, +} + +/// Agent management commands. +#[derive(Subcommand, Debug, Clone)] +pub enum AgentCommand { + /// List available agents. + #[command(alias = "ls")] + List, +} + +/// Command group for workspace management. +#[derive(Parser, Debug, Clone)] +pub struct WorkspaceCommandGroup { + #[command(subcommand)] + pub command: WorkspaceCommand, +} + +#[derive(Subcommand, Debug, Clone)] +pub enum WorkspaceCommand { + /// Synchronize a directory for semantic search. + Sync { + /// Path to the directory to sync + #[arg(default_value = ".")] + path: PathBuf, + + /// Automatically initialize the workspace before syncing if it has not + /// been initialized yet. + #[arg(long)] + init: bool, + }, + /// List all workspaces. + List { + /// Output in machine-readable format + #[arg(short, long)] + porcelain: bool, + }, + + /// Query the workspace. + Query { + /// Search query. + query: String, + + /// Path to the directory to index (used when no subcommand is + /// provided). + #[arg(default_value = ".")] + path: PathBuf, + + /// Maximum number of results to return. + #[arg(short, long, default_value = "10")] + limit: usize, + + /// Number of highest probability tokens to consider (1-1000). + #[arg(long)] + top_k: Option, + + /// Describe your intent or goal to filter results for relevance. + #[arg(long, short = 'r')] + use_case: String, + + /// Filter results to files starting with this prefix. + #[arg(long)] + starts_with: Option, + + /// Filter results to files ending with this suffix. + #[arg(long)] + ends_with: Option, + }, + + /// Show workspace information for an indexed directory. + Info { + /// Path to the directory to get information for + #[arg(default_value = ".")] + path: PathBuf, + }, + + /// Delete one or more workspaces. + Delete { + /// Workspace IDs to delete + workspace_ids: Vec, + }, + + /// Show sync status of all files in the workspace. + Status { + /// Path to the directory to check status for + #[arg(default_value = ".")] + path: PathBuf, + + /// Output in machine-readable format + #[arg(short, long)] + porcelain: bool, + }, + + /// Initialize an empty workspace for the provided directory + Init { + /// Path to the directory to initialize as a workspace + #[arg(default_value = ".")] + path: PathBuf, + + /// Automatically confirm initialization without prompting + #[arg(short = 'y', long)] + yes: bool, + }, +} + +/// Command group for listing resources. +#[derive(Parser, Debug, Clone)] +pub struct ListCommandGroup { + #[command(subcommand)] + pub command: ListCommand, + + /// Output in machine-readable format. + #[arg(long, global = true)] + pub porcelain: bool, +} + +#[derive(Subcommand, Debug, Clone)] +pub enum ListCommand { + /// List available agents. + #[command(alias = "agents")] + Agent { + /// Shows only custom agents + #[arg(long)] + custom: bool, + }, + + /// List available API providers. + #[command(alias = "providers")] + Provider { + /// Filter providers by type (e.g., llm, context_engine). Can be + /// specified multiple times. + #[arg(long = "type", short = 't')] + types: Vec, + }, + + /// List available models. + #[command(alias = "models")] + Model, + + /// List available commands. + #[command(hide = true, alias = "commands")] + Command { + /// Shows only custom commands + #[arg(long)] + custom: bool, + }, + + /// List configuration values. + #[command(alias = "configs")] + Config, + + /// List tools for a specific agent. + #[command(alias = "tools")] + Tool { + /// Agent ID to list tools for. + agent: AgentId, + }, + + /// List MCP servers. + #[command(alias = "mcps")] + Mcp, + + /// List conversation history. + #[command(alias = "session")] + Conversation { + /// Show child conversations of a parent conversation. + #[arg(long)] + parent: Option, + }, + + /// List custom commands. + #[command(alias = "cmds")] + Cmd, + + /// List available skills. + #[command(alias = "skills")] + Skill { + /// Shows only custom skills + #[arg(long)] + custom: bool, + }, + + /// List files and directories in the current workspace. + /// + /// Includes hidden files and directories (dotfiles), respects .gitignore, + /// and outputs one path per line. Directories are suffixed with `/`. + #[command(alias = "files")] + File, +} + +/// Shell extension commands. +#[derive(Subcommand, Debug, Clone)] +pub enum ZshCommandGroup { + /// Generate shell plugin script + Plugin, + /// Generate shell theme + Theme, + /// Run diagnostics on shell environment + Doctor, + + /// Get rprompt information (model and conversation stats) for shell + /// integration. + Rprompt, + + /// Setup zsh integration by updating .zshrc with plugin and theme + Setup, + + /// Show keyboard shortcuts for ZSH line editor + Keyboard, + + /// Format buffer text by wrapping file paths in @[...] syntax. + /// + /// Used by the zsh plugin to delegate path detection and wrapping to + /// Rust where the logic is well-tested across all terminal environments. + Format { + /// The text buffer to format. + #[arg(long)] + buffer: String, + }, +} + +/// Command group for MCP server management. +#[derive(Parser, Debug, Clone)] +pub struct McpCommandGroup { + #[command(subcommand)] + pub command: McpCommand, + + /// Output in machine-readable format. + #[arg(long, global = true)] + pub porcelain: bool, +} + +#[derive(Subcommand, Debug, Clone)] +pub enum McpCommand { + /// Import server configuration from JSON. + Import(McpImportArgs), + + /// List configured servers. + List, + + /// Remove a configured server. + Remove(McpRemoveArgs), + + /// Show server configuration details. + Show(McpShowArgs), + + /// Reload servers and rebuild caches. + Reload, + + /// Authenticate with an OAuth-enabled MCP server. + Login(McpAuthArgs), + + /// Remove stored OAuth credentials for an MCP server. + Logout(McpLogoutArgs), +} + +#[derive(Parser, Debug, Clone)] +pub struct McpImportArgs { + /// JSON configuration to import. + #[arg()] + pub json: String, + + /// Configuration scope. + #[arg(short = 's', long = "scope", default_value = "local")] + pub scope: Scope, +} + +#[derive(Parser, Debug, Clone)] +pub struct McpRemoveArgs { + /// Configuration scope. + #[arg(short = 's', long = "scope", default_value = "local")] + pub scope: Scope, + + /// Name of the server to remove. + pub name: String, +} + +#[derive(Parser, Debug, Clone)] +pub struct McpShowArgs { + /// Name of the server to show details for. + pub name: String, +} + +#[derive(Parser, Debug, Clone)] +pub struct McpAuthArgs { + /// Name of the MCP server to authenticate with. + pub name: String, +} + +#[derive(Parser, Debug, Clone)] +pub struct McpLogoutArgs { + /// Name of the MCP server to remove credentials for, or "all" to + /// remove all MCP OAuth credentials. + pub name: String, +} + +/// Configuration scope for settings. +#[derive(Copy, Clone, Debug, ValueEnum, Default)] +pub enum Scope { + /// Local configuration (project-specific). + #[default] + Local, + /// User configuration (global to the user). + User, +} + +impl From for forge_domain::Scope { + fn from(value: Scope) -> Self { + match value { + Scope::Local => forge_domain::Scope::Local, + Scope::User => forge_domain::Scope::User, + } + } +} + +/// Transport protocol for communication. +#[derive(Copy, Clone, Debug, ValueEnum)] +#[clap(rename_all = "lower")] +pub enum Transport { + /// Standard input/output communication. + Stdio, + /// Server-sent events communication. + Sse, +} + +/// Command group for configuration management. +#[derive(Parser, Debug, Clone)] +pub struct ConfigCommandGroup { + #[command(subcommand)] + pub command: ConfigCommand, + + /// Output in machine-readable format. + #[arg(long, global = true)] + pub porcelain: bool, +} + +#[derive(Subcommand, Debug, Clone)] +pub enum ConfigCommand { + /// Set a configuration value. + Set(ConfigSetArgs), + + /// Get a configuration value. + Get(ConfigGetArgs), + + /// List configuration values. + List, + + /// Print the path to the global config file. + Path, + + /// Migrate the legacy ~/forge directory to ~/.forge. + Migrate, +} + +/// Arguments for `forge config set`. +#[derive(Parser, Debug, Clone)] +pub struct ConfigSetArgs { + #[command(subcommand)] + pub field: ConfigSetField, +} + +/// Arguments for `forge config get`. +#[derive(Parser, Debug, Clone)] +pub struct ConfigGetArgs { + #[command(subcommand)] + pub field: ConfigGetField, +} + +/// Type-safe subcommands for `forge config set`. +#[derive(Subcommand, Debug, Clone)] +pub enum ConfigSetField { + /// Set the active model and provider atomically. + Model { + /// Provider ID to set as default. + provider: ProviderId, + /// Model ID to set as default. + model: ModelId, + }, + /// Set the provider and model for commit message generation. + Commit { + /// Provider ID to use for commit message generation. + provider: ProviderId, + /// Model ID to use for commit message generation. + model: ModelId, + }, + /// Set the provider and model for command suggestion generation. + Suggest { + /// Provider ID to use for command suggestion generation. + provider: ProviderId, + /// Model ID to use for command suggestion generation. + model: ModelId, + }, + /// Set the reasoning effort level applied to all agents. + ReasoningEffort { + /// Effort level: none, minimal, low, medium, high, xhigh, max. + effort: Effort, + }, +} + +/// Type-safe subcommands for `forge config get`. +#[derive(Subcommand, Debug, Clone)] +pub enum ConfigGetField { + /// Get the active model. + Model, + /// Get the active provider. + Provider, + /// Get the commit message generation config. + Commit, + /// Get the command suggestion generation config. + Suggest, + /// Get the reasoning effort level. + ReasoningEffort, +} + +/// Command group for conversation management. +#[derive(Parser, Debug, Clone)] +pub struct ConversationCommandGroup { + #[command(subcommand)] + pub command: ConversationCommand, +} +#[derive(Subcommand, Debug, Clone)] +pub enum ConversationCommand { + /// List conversation history. + List { + /// Output in machine-readable format. + #[arg(long)] + porcelain: bool, + }, + + /// Create a new conversation. + New, + + /// Export conversation as JSON or HTML. + Dump { + /// Conversation ID to export. + id: ConversationId, + + /// Export as HTML instead of JSON. + #[arg(long)] + html: bool, + }, + + /// Compact conversation to reduce token usage. + Compact { + /// Conversation ID to compact. + id: ConversationId, + }, + + /// Retry last command without modifying context. + Retry { + /// Conversation ID to retry. + id: ConversationId, + }, + + /// Resume conversation in interactive mode. + Resume { + /// Conversation ID to resume. + id: ConversationId, + }, + + /// Show last assistant message. + Show { + /// Conversation ID. + id: ConversationId, + + /// Print raw markdown without rendering. + #[arg(long)] + md: bool, + }, + + /// Show conversation details. + Info { + /// Conversation ID. + id: ConversationId, + }, + + /// Show conversation statistics. + Stats { + /// Conversation ID. + id: ConversationId, + + /// Output in machine-readable format. + #[arg(long)] + porcelain: bool, + }, + + /// Clone conversation with a new ID. + Clone { + /// Conversation ID to clone. + id: ConversationId, + + /// Output in machine-readable format. + #[arg(long)] + porcelain: bool, + }, + + /// Delete a conversation permanently. + Delete { + /// Conversation ID to delete. + id: String, + }, + + /// Rename a conversation. + Rename { + /// Conversation ID to rename. + id: ConversationId, + + /// New name for the conversation. + name: String, + }, +} + +/// Command group for provider authentication management. +#[derive(Parser, Debug, Clone)] +pub struct ProviderCommandGroup { + #[command(subcommand)] + pub command: ProviderCommand, + + /// Output in machine-readable format. + #[arg(long, global = true)] + pub porcelain: bool, +} + +#[derive(Subcommand, Debug, Clone)] +pub enum ProviderCommand { + /// Authenticate with an API provider. + /// + /// Shows an interactive menu when no provider name is specified. + Login { + /// Provider name to authenticate with. + provider: Option, + }, + + /// Remove provider credentials. + /// + /// Shows an interactive menu when no provider name is specified. + Logout { + /// Provider name to log out from. + provider: Option, + }, + + /// List available providers. + List { + /// Filter providers by type (e.g., llm, context_engine). Can be + /// specified multiple times. + #[arg(long = "type", short = 't')] + types: Vec, + }, +} + +/// Group of Commit-related commands +#[derive(Parser, Debug, Clone)] +pub struct CommitCommandGroup { + /// Preview the commit message without committing + #[arg(long)] + pub preview: bool, + + /// Maximum git diff size in bytes (default: 100k) + /// + /// Limits the size of the git diff sent to the AI model. Large diffs are + /// truncated to save tokens and reduce API costs. Minimum value is 5000 + /// bytes. + #[arg(long = "max-diff", default_value = "100000", value_parser = clap::builder::RangedI64ValueParser::::new().range(5000..))] + pub max_diff_size: Option, + + /// Git diff content (used internally for piped input) + /// + /// This field is populated when diff content is piped to the commit + /// command. Users typically don't set this directly; instead, they pipe + /// diff content: `git diff | forge commit --preview` + #[arg(skip)] + pub diff: Option, + + /// Additional text to customize the commit message + /// + /// Provide additional context or instructions for the AI to use when + /// generating the commit message. Multiple words can be provided without + /// quotes: `forge commit fix typo in readme` + pub text: Vec, +} + +/// Group of Data-related commands +#[derive(Parser, Debug, Clone)] +pub struct DataCommandGroup { + /// Path to JSONL file to process + #[arg(long)] + pub input: String, + + /// Path to JSON schema file for LLM tool definition + #[arg(long)] + pub schema: String, + + /// Path to Handlebars template file for system prompt + #[arg(long)] + pub system_prompt: Option, + + /// Path to Handlebars template file for user prompt + #[arg(long)] + pub user_prompt: Option, + + /// Maximum number of concurrent LLM requests + #[arg(long, default_value = "10")] + pub concurrency: usize, +} + +impl From for forge_domain::DataGenerationParameters { + fn from(value: DataCommandGroup) -> Self { + Self { + input: value.input.into(), + schema: value.schema.into(), + system_prompt: value.system_prompt.map(Into::into), + user_prompt: value.user_prompt.map(Into::into), + concurrency: value.concurrency, + } + } +} + +/// VS Code integration commands. +#[derive(Subcommand, Debug, Clone)] +pub enum VscodeCommand { + /// Install the Forge VS Code extension. + InstallExtension, +} + +/// Update command arguments. +#[derive(Parser, Debug, Clone)] +pub struct UpdateArgs { + /// Skip the confirmation prompt when applying updates. + #[arg(long, default_value_t = false)] + pub no_confirm: bool, +} + +/// Arguments for the `forge logs` command. +#[derive(Parser, Debug, Clone)] +pub struct LogsArgs { + /// Number of lines to show from the end of the log file. + #[arg(long, short = 'n', default_value_t = 20)] + pub lines: usize, + + /// Do not follow the log output; exit after printing the last lines. + #[arg(long)] + pub no_follow: bool, + + /// List all available log files instead of tailing one. + #[arg(long, short = 'l')] + pub list: bool, + + /// Path to a specific log file to tail. Defaults to the most recent log + /// file in the forge logs directory. + #[arg(long, short = 'f')] + pub file: Option, +} + +#[cfg(test)] +mod tests { + use clap::Parser; + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_data_command_group_conversion() { + use std::path::PathBuf; + + let fixture = DataCommandGroup { + input: "path/to/input.jsonl".to_string(), + schema: "path/to/schema.json".to_string(), + system_prompt: Some("system prompt".to_string()), + user_prompt: None, + concurrency: 5, + }; + let actual: forge_domain::DataGenerationParameters = fixture.into(); + let expected = forge_domain::DataGenerationParameters { + input: PathBuf::from("path/to/input.jsonl"), + schema: PathBuf::from("path/to/schema.json"), + system_prompt: Some(PathBuf::from("system prompt")), + user_prompt: None, + concurrency: 5, + }; + assert_eq!(actual, expected); + } + + #[test] + fn test_commit_default_max_diff_size() { + let fixture = Cli::parse_from(["forge", "commit", "--preview"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Commit(commit)) => commit.max_diff_size, + _ => panic!("Expected Commit command"), + }; + let expected = Some(100000); + assert_eq!(actual, expected); + } + + #[test] + fn test_commit_custom_max_diff_size() { + let fixture = Cli::parse_from(["forge", "commit", "--preview", "--max-diff", "50000"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Commit(commit)) => commit.max_diff_size, + _ => panic!("Expected Commit command"), + }; + let expected = Some(50000); + assert_eq!(actual, expected); + } + + #[test] + fn test_config_set_with_provider_and_model() { + let fixture = Cli::parse_from([ + "forge", + "config", + "set", + "model", + "anthropic", + "claude-sonnet-4-20250514", + ]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Config(config)) => match config.command { + ConfigCommand::Set(args) => match args.field { + ConfigSetField::Model { provider, model } => { + Some((provider.to_string(), model.as_str().to_string())) + } + _ => None, + }, + _ => None, + }, + _ => None, + }; + let expected = Some(( + "Anthropic".to_string(), + "claude-sonnet-4-20250514".to_string(), + )); + assert_eq!(actual, expected); + } + + #[test] + fn test_config_list() { + let fixture = Cli::parse_from(["forge", "config", "list"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Config(config)) => matches!(config.command, ConfigCommand::List), + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_config_get_specific_field() { + let fixture = Cli::parse_from(["forge", "config", "get", "model"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Config(config)) => match config.command { + ConfigCommand::Get(args) => matches!(args.field, ConfigGetField::Model), + _ => panic!("Expected ConfigCommand::Get"), + }, + _ => panic!("Expected TopLevelCommand::Config"), + }; + assert!(actual); + } + + #[test] + fn test_config_set_commit_with_provider_and_model() { + let fixture = Cli::parse_from([ + "forge", + "config", + "set", + "commit", + "anthropic", + "claude-haiku-4-20250514", + ]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Config(config)) => match config.command { + ConfigCommand::Set(args) => match args.field { + ConfigSetField::Commit { provider, model } => { + Some((provider.to_string(), model.as_str().to_string())) + } + _ => None, + }, + _ => None, + }, + _ => None, + }; + let expected = Some(( + "Anthropic".to_string(), + "claude-haiku-4-20250514".to_string(), + )); + assert_eq!(actual, expected); + } + + #[test] + fn test_conversation_list() { + let fixture = Cli::parse_from(["forge", "conversation", "list"]); + let is_list = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => { + matches!(conversation.command, ConversationCommand::List { .. }) + } + _ => false, + }; + assert_eq!(is_list, true); + } + + #[test] + fn test_session_alias_list() { + let fixture = Cli::parse_from(["forge", "session", "list"]); + let is_list = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => { + matches!(conversation.command, ConversationCommand::List { .. }) + } + _ => false, + }; + assert_eq!(is_list, true); + } + + #[test] + fn test_agent_id_long_flag() { + let fixture = Cli::parse_from(["forge", "--agent", "sage"]); + assert_eq!(fixture.agent, Some(AgentId::new("sage"))); + } + + #[test] + fn test_agent_id_short_alias() { + let fixture = Cli::parse_from(["forge", "--aid", "muse"]); + assert_eq!(fixture.agent, Some(AgentId::new("muse"))); + } + + #[test] + fn test_agent_id_with_prompt() { + let fixture = Cli::parse_from(["forge", "--agent", "forge", "-p", "test prompt"]); + assert_eq!(fixture.agent, Some(AgentId::new("forge"))); + assert_eq!(fixture.prompt, Some("test prompt".to_string())); + } + + #[test] + fn test_agent_id_not_provided() { + let fixture = Cli::parse_from(["forge"]); + assert_eq!(fixture.agent, None); + } + + #[test] + fn test_conversation_dump_json_with_id() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "dump", + "550e8400-e29b-41d4-a716-446655440000", + ]); + let (id, html) = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Dump { id, html } => (id, html), + _ => (ConversationId::default(), true), + }, + _ => (ConversationId::default(), true), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440000").unwrap() + ); + assert_eq!(html, false); // JSON is default + } + + #[test] + fn test_conversation_dump_html_with_id() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "dump", + "550e8400-e29b-41d4-a716-446655440001", + "--html", + ]); + let (id, html) = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Dump { id, html } => (id, html), + _ => (ConversationId::default(), false), + }, + _ => (ConversationId::default(), false), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440001").unwrap() + ); + assert_eq!(html, true); + } + + #[test] + fn test_conversation_retry_with_id() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "retry", + "550e8400-e29b-41d4-a716-446655440002", + ]); + let id = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Retry { id } => id, + _ => ConversationId::default(), + }, + _ => ConversationId::default(), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440002").unwrap() + ); + } + + #[test] + fn test_conversation_compact_with_id() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "compact", + "550e8400-e29b-41d4-a716-446655440003", + ]); + let id = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Compact { id } => id, + _ => ConversationId::default(), + }, + _ => ConversationId::default(), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440003").unwrap() + ); + } + + #[test] + fn test_conversation_last_with_id() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "show", + "550e8400-e29b-41d4-a716-446655440004", + ]); + let (id, md) = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Show { id, md } => (id, md), + _ => (ConversationId::default(), false), + }, + _ => (ConversationId::default(), false), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440004").unwrap() + ); + assert_eq!(md, false); + } + + #[test] + fn test_conversation_show_with_md_flag() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "show", + "550e8400-e29b-41d4-a716-446655440004", + "--md", + ]); + let (id, md) = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Show { id, md } => (id, md), + _ => (ConversationId::default(), false), + }, + _ => (ConversationId::default(), false), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440004").unwrap() + ); + assert_eq!(md, true); + } + + #[test] + fn test_conversation_resume() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "resume", + "550e8400-e29b-41d4-a716-446655440005", + ]); + let id = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Resume { id } => id, + _ => ConversationId::default(), + }, + _ => ConversationId::default(), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440005").unwrap() + ); + } + + #[test] + fn test_list_tools_command_with_agent() { + let fixture = Cli::parse_from(["forge", "list", "tool", "sage"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => match list.command { + ListCommand::Tool { agent } => agent, + _ => AgentId::default(), + }, + _ => AgentId::default(), + }; + let expected = AgentId::new("sage"); + assert_eq!(actual, expected); + } + + #[test] + fn test_list_conversation_command() { + let fixture = Cli::parse_from(["forge", "list", "conversation"]); + let is_conversation_list = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => { + matches!(list.command, ListCommand::Conversation { .. }) + } + _ => false, + }; + assert_eq!(is_conversation_list, true); + } + + #[test] + fn test_list_session_alias_command() { + let fixture = Cli::parse_from(["forge", "list", "session"]); + let is_conversation_list = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => { + matches!(list.command, ListCommand::Conversation { .. }) + } + _ => false, + }; + assert_eq!(is_conversation_list, true); + } + + #[test] + fn test_list_command_without_custom_flag() { + let fixture = Cli::parse_from(["forge", "list", "command"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => match list.command { + ListCommand::Command { custom } => custom, + _ => true, + }, + _ => true, + }; + let expected = false; + assert_eq!(actual, expected); + } + + #[test] + fn test_list_command_with_custom_flag() { + let fixture = Cli::parse_from(["forge", "list", "command", "--custom"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => match list.command { + ListCommand::Command { custom } => custom, + _ => false, + }, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_cmd_list_with_custom_flag() { + let fixture = Cli::parse_from(["forge", "cmd", "list", "--custom"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Cmd(cmd_group)) => match cmd_group.command { + CmdCommand::List { custom } => custom, + _ => false, + }, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_command_list_with_custom_flag() { + let fixture = Cli::parse_from(["forge", "command", "list", "--custom"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Cmd(cmd_group)) => match cmd_group.command { + CmdCommand::List { custom } => custom, + _ => false, + }, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_info_command_without_porcelain() { + let fixture = Cli::parse_from(["forge", "info"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Info { porcelain, .. }) => porcelain, + _ => true, + }; + let expected = false; + assert_eq!(actual, expected); + } + + #[test] + fn test_info_command_with_porcelain() { + let fixture = Cli::parse_from(["forge", "info", "--porcelain"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Info { porcelain, .. }) => porcelain, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_info_command_with_conversation_id() { + let fixture = Cli::parse_from([ + "forge", + "info", + "--conversation-id", + "550e8400-e29b-41d4-a716-446655440006", + ]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Info { conversation_id, .. }) => conversation_id, + _ => None, + }; + let expected = Some(ConversationId::parse("550e8400-e29b-41d4-a716-446655440006").unwrap()); + assert_eq!(actual, expected); + } + + #[test] + fn test_info_command_with_cid_alias() { + let fixture = Cli::parse_from([ + "forge", + "info", + "--cid", + "550e8400-e29b-41d4-a716-446655440007", + ]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Info { conversation_id, .. }) => conversation_id, + _ => None, + }; + let expected = Some(ConversationId::parse("550e8400-e29b-41d4-a716-446655440007").unwrap()); + assert_eq!(actual, expected); + } + + #[test] + fn test_info_command_with_conversation_id_and_porcelain() { + let fixture = Cli::parse_from([ + "forge", + "info", + "--cid", + "550e8400-e29b-41d4-a716-446655440008", + "--porcelain", + ]); + let (conversation_id, porcelain) = match fixture.subcommands { + Some(TopLevelCommand::Info { conversation_id, porcelain }) => { + (conversation_id, porcelain) + } + _ => (None, false), + }; + assert_eq!( + conversation_id, + Some(ConversationId::parse("550e8400-e29b-41d4-a716-446655440008").unwrap()) + ); + assert_eq!(porcelain, true); + } + + #[test] + fn test_list_agents_without_porcelain() { + let fixture = Cli::parse_from(["forge", "list", "agents"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => list.porcelain, + _ => true, + }; + let expected = false; + assert_eq!(actual, expected); + } + + #[test] + fn test_list_agents_with_porcelain() { + let fixture = Cli::parse_from(["forge", "list", "agents", "--porcelain"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => list.porcelain, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_mcp_list_with_porcelain() { + let fixture = Cli::parse_from(["forge", "mcp", "list", "--porcelain"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Mcp(mcp)) => mcp.porcelain, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_conversation_list_with_porcelain() { + let fixture = Cli::parse_from(["forge", "conversation", "list", "--porcelain"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::List { porcelain } => porcelain, + _ => false, + }, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_list_models_with_porcelain() { + let fixture = Cli::parse_from(["forge", "list", "models", "--porcelain"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => list.porcelain, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_config_list_with_porcelain() { + let fixture = Cli::parse_from(["forge", "config", "list", "--porcelain"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Config(config)) => config.porcelain, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_conversation_info_with_id() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "info", + "550e8400-e29b-41d4-a716-446655440009", + ]); + let id = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Info { id } => id, + _ => ConversationId::default(), + }, + _ => ConversationId::default(), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440009").unwrap() + ); + } + + #[test] + fn test_conversation_stats_with_porcelain() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "stats", + "550e8400-e29b-41d4-a716-446655440010", + "--porcelain", + ]); + let (id, porcelain) = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Stats { id, porcelain } => (id, porcelain), + _ => (ConversationId::default(), false), + }, + _ => (ConversationId::default(), false), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440010").unwrap() + ); + assert_eq!(porcelain, true); + } + + #[test] + fn test_prompt_command() { + let fixture = Cli::parse_from(["forge", "zsh", "rprompt"]); + let r_prompt = matches!( + fixture.subcommands, + Some(TopLevelCommand::Zsh(ZshCommandGroup::Rprompt)) + ); + assert!(r_prompt); + } + + #[test] + fn test_session_alias_dump() { + let fixture = Cli::parse_from([ + "forge", + "session", + "dump", + "550e8400-e29b-41d4-a716-446655440011", + ]); + let (id, html) = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Dump { id, html } => (id, html), + _ => (ConversationId::default(), true), + }, + _ => (ConversationId::default(), true), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440011").unwrap() + ); + assert_eq!(html, false); + } + + #[test] + fn test_session_alias_retry() { + let fixture = Cli::parse_from([ + "forge", + "session", + "retry", + "550e8400-e29b-41d4-a716-446655440012", + ]); + let id = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Retry { id } => id, + _ => ConversationId::default(), + }, + _ => ConversationId::default(), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440012").unwrap() + ); + } + + #[test] + fn test_prompt_with_conversation_id() { + let fixture = Cli::parse_from([ + "forge", + "-p", + "hello", + "--conversation-id", + "550e8400-e29b-41d4-a716-446655440000", + ]); + let actual = fixture.conversation_id; + let expected = Some(ConversationId::parse("550e8400-e29b-41d4-a716-446655440000").unwrap()); + assert_eq!(actual, expected); + } + + #[test] + fn test_conversation_id_without_prompt() { + let fixture = Cli::parse_from([ + "forge", + "--conversation-id", + "550e8400-e29b-41d4-a716-446655440000", + ]); + let actual = fixture.conversation_id; + let expected = Some(ConversationId::parse("550e8400-e29b-41d4-a716-446655440000").unwrap()); + assert_eq!(actual, expected); + } + + #[test] + fn test_conversation_clone_with_id() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "clone", + "550e8400-e29b-41d4-a716-446655440013", + ]); + let id = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Clone { id, .. } => id, + _ => ConversationId::default(), + }, + _ => ConversationId::default(), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440013").unwrap() + ); + } + + #[test] + fn test_conversation_clone_with_porcelain() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "clone", + "550e8400-e29b-41d4-a716-446655440014", + "--porcelain", + ]); + let (id, porcelain) = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => match conversation.command { + ConversationCommand::Clone { id, porcelain } => (id, porcelain), + _ => (ConversationId::default(), false), + }, + _ => (ConversationId::default(), false), + }; + assert_eq!( + id, + ConversationId::parse("550e8400-e29b-41d4-a716-446655440014").unwrap() + ); + assert_eq!(porcelain, true); + } + + #[test] + fn test_cmd_command_with_args() { + let fixture = + Cli::parse_from(["forge", "cmd", "execute", "custom-command", "arg1", "arg2"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Cmd(run_group)) => match run_group.command { + CmdCommand::Execute { commands } => commands.join(" "), + _ => panic!("Expected Execute command"), + }, + _ => panic!("Expected Cmd command"), + }; + let expected = "custom-command arg1 arg2".to_string(); + assert_eq!(actual, expected); + } + + #[test] + fn test_is_interactive_without_flags() { + let fixture = Cli::parse_from(["forge"]); + let actual = fixture.is_interactive(); + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_commit_with_custom_text() { + let fixture = Cli::parse_from(["forge", "commit", "fix", "typo", "in", "readme"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Commit(commit)) => commit.text, + _ => panic!("Expected Commit command"), + }; + let expected = ["fix", "typo", "in", "readme"] + .iter() + .map(|s| s.to_string()) + .collect::>(); + assert_eq!(actual, expected); + } + + #[test] + fn test_commit_without_custom_text() { + let fixture = Cli::parse_from(["forge", "commit", "--preview"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Commit(commit)) => commit.text, + _ => panic!("Expected Commit command"), + }; + let expected: Vec = vec![]; + assert_eq!(actual, expected); + } + + #[test] + fn test_commit_with_text_and_flags() { + let fixture = Cli::parse_from([ + "forge", + "commit", + "--preview", + "--max-diff", + "50000", + "update", + "docs", + ]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Commit(commit)) => { + (commit.preview, commit.max_diff_size, commit.text) + } + _ => panic!("Expected Commit command"), + }; + let expected = ( + true, + Some(50000), + ["update", "docs"] + .iter() + .map(|s| s.to_string()) + .collect::>(), + ); + assert_eq!(actual, expected); + } + + #[test] + fn test_list_skill_command() { + let fixture = Cli::parse_from(["forge", "list", "skill"]); + let is_skill_list = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => matches!(list.command, ListCommand::Skill { .. }), + _ => false, + }; + assert_eq!(is_skill_list, true); + } + + #[test] + fn test_list_skills_alias_command() { + let fixture = Cli::parse_from(["forge", "list", "skills"]); + let is_skill_list = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => matches!(list.command, ListCommand::Skill { .. }), + _ => false, + }; + assert_eq!(is_skill_list, true); + } + + #[test] + fn test_conversation_delete_with_id() { + let fixture = Cli::parse_from([ + "forge", + "conversation", + "delete", + "550e8400-e29b-41d4-a716-446655440000", + ]); + let is_delete_with_id = match fixture.subcommands { + Some(TopLevelCommand::Conversation(conversation)) => { + matches!(conversation.command, ConversationCommand::Delete { id: _ }) + } + _ => false, + }; + assert_eq!(is_delete_with_id, true); + } + + #[test] + fn test_list_skill_with_porcelain() { + let fixture = Cli::parse_from(["forge", "list", "skill", "--porcelain"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => list.porcelain, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_prompt_with_leading_hyphen() { + let fixture = Cli::parse_from(["forge", "-p", "- hi"]); + assert_eq!(fixture.prompt, Some("- hi".to_string())); + } + + #[test] + fn test_prompt_with_hyphen_flag_like_value() { + let fixture = Cli::parse_from(["forge", "-p", "-test"]); + assert_eq!(fixture.prompt, Some("-test".to_string())); + } + + #[test] + fn test_prompt_with_double_hyphen() { + let fixture = Cli::parse_from(["forge", "-p", "--something"]); + assert_eq!(fixture.prompt, Some("--something".to_string())); + } + + #[test] + fn test_suggest_with_dash_prefixed_prompt() { + let fixture = Cli::parse_from(["forge", "suggest", "--- date"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Suggest { prompt }) => prompt, + _ => panic!("Expected suggest subcommand"), + }; + let expected = "--- date".to_string(); + assert_eq!(actual, expected); + } + + #[test] + fn test_suggest_with_double_dash_prompt() { + let fixture = Cli::parse_from(["forge", "suggest", "--date tomorrow"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Suggest { prompt }) => prompt, + _ => panic!("Expected suggest subcommand"), + }; + let expected = "--date tomorrow".to_string(); + assert_eq!(actual, expected); + } + + #[test] + fn test_suggest_with_single_dash_prompt() { + let fixture = Cli::parse_from(["forge", "suggest", "-v file.txt"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Suggest { prompt }) => prompt, + _ => panic!("Expected suggest subcommand"), + }; + let expected = "-v file.txt".to_string(); + assert_eq!(actual, expected); + } + + #[test] + fn test_terminal_theme_zsh() { + let fixture = Cli::parse_from(["forge", "zsh", "theme"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Zsh(terminal)) => { + matches!(terminal, ZshCommandGroup::Theme) + } + _ => false, + }; + assert_eq!(actual, true); + } + + #[test] + fn test_terminal_plugin_zsh() { + let fixture = Cli::parse_from(["forge", "zsh", "plugin"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Zsh(terminal)) => { + matches!(terminal, ZshCommandGroup::Plugin) + } + _ => false, + }; + assert_eq!(actual, true); + } + + #[test] + fn test_zsh_doctor() { + let fixture = Cli::parse_from(["forge", "zsh", "doctor"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Zsh(terminal)) => { + matches!(terminal, ZshCommandGroup::Doctor) + } + _ => false, + }; + assert_eq!(actual, true); + } + + #[test] + fn test_zsh_setup() { + let fixture = Cli::parse_from(["forge", "zsh", "setup"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Zsh(terminal)) => { + matches!(terminal, ZshCommandGroup::Setup) + } + _ => false, + }; + assert_eq!(actual, true); + } + + #[test] + fn test_zsh_keyboard() { + let fixture = Cli::parse_from(["forge", "zsh", "keyboard"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Zsh(terminal)) => { + matches!(terminal, ZshCommandGroup::Keyboard) + } + _ => false, + }; + assert_eq!(actual, true); + } + + #[test] + fn test_zsh_format() { + let fixture = Cli::parse_from(["forge", "zsh", "format", "--buffer", "hello world"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Zsh(ZshCommandGroup::Format { buffer })) => { + buffer == "hello world" + } + _ => false, + }; + assert_eq!(actual, true); + } + + #[test] + fn test_setup_alias() { + let fixture = Cli::parse_from(["forge", "setup"]); + let actual = matches!(fixture.subcommands, Some(TopLevelCommand::Setup)); + assert_eq!(actual, true); + } + + #[test] + fn test_doctor_alias() { + let fixture = Cli::parse_from(["forge", "doctor"]); + let actual = matches!(fixture.subcommands, Some(TopLevelCommand::Doctor)); + assert_eq!(actual, true); + } + + #[test] + fn test_install_vscode_extension() { + let fixture = Cli::parse_from(["forge", "vscode", "install-extension"]); + let actual = matches!( + fixture.subcommands, + Some(TopLevelCommand::Vscode(VscodeCommand::InstallExtension)) + ); + assert_eq!(actual, true); + } + + #[test] + fn test_list_agent_with_custom_flag() { + let fixture = Cli::parse_from(["forge", "list", "agent", "--custom"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => match list.command { + ListCommand::Agent { custom } => custom, + _ => false, + }, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_list_agent_without_custom_flag() { + let fixture = Cli::parse_from(["forge", "list", "agent"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => match list.command { + ListCommand::Agent { custom } => custom, + _ => true, + }, + _ => true, + }; + let expected = false; + assert_eq!(actual, expected); + } + + #[test] + fn test_list_skill_with_custom_flag() { + let fixture = Cli::parse_from(["forge", "list", "skill", "--custom"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => match list.command { + ListCommand::Skill { custom } => custom, + _ => false, + }, + _ => false, + }; + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn test_list_skill_without_custom_flag() { + let fixture = Cli::parse_from(["forge", "list", "skill"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::List(list)) => match list.command { + ListCommand::Skill { custom } => custom, + _ => true, + }, + _ => true, + }; + let expected = false; + assert_eq!(actual, expected); + } + + #[test] + fn test_update_with_no_confirm() { + let fixture = Cli::parse_from(["forge", "update", "--no-confirm"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Update(args)) => args.no_confirm, + _ => panic!("Expected Update command"), + }; + assert!(actual); + } + + #[test] + fn test_update_without_no_confirm() { + let fixture = Cli::parse_from(["forge", "update"]); + let actual = match fixture.subcommands { + Some(TopLevelCommand::Update(args)) => args.no_confirm, + _ => panic!("Expected Update command"), + }; + assert!(!actual); + } +} diff --git a/crates/forge_main/src/completer/command.rs b/crates/forge_main/src/completer/command.rs new file mode 100644 index 0000000000000000000000000000000000000000..04daab5ec37316e5b609d3e2a0eca063cb7a3de2 --- /dev/null +++ b/crates/forge_main/src/completer/command.rs @@ -0,0 +1,87 @@ +use std::sync::Arc; + +use forge_select::ForgeWidget; + +use crate::completer::input_completer::InputSuggestion; +use crate::completer::search_term::Span; +use crate::model::{ForgeCommand, ForgeCommandManager}; + +/// A display wrapper for `ForgeCommand` that renders the name and description +/// side-by-side for the interactive picker. +#[derive(Clone)] +struct CommandRow(ForgeCommand); + +impl std::fmt::Display for CommandRow { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:<30} {}", self.0.name, self.0.description) + } +} + +#[derive(Clone)] +pub struct CommandCompleter(Arc); + +impl CommandCompleter { + pub fn new(command_manager: Arc) -> Self { + Self(command_manager) + } +} + +impl CommandCompleter { + pub fn complete(&mut self, line: &str, _: usize) -> Vec { + // Determine which sentinel the user typed (`:` or `/`), defaulting to `/`. + let sentinel = if line.starts_with(':') { ':' } else { '/' }; + + // Build the list of display names using the same sentinel the user typed. + let commands: Vec = self + .0 + .list() + .into_iter() + .filter_map(|cmd| { + let display_name = if cmd.name.starts_with('!') { + cmd.name.clone() + } else { + format!("{}{}", sentinel, cmd.name) + }; + + // Only include commands that match what the user has typed so far. + if display_name.starts_with(line) { + Some(CommandRow(ForgeCommand { + name: display_name, + description: cmd.description, + value: cmd.value, + })) + } else { + None + } + }) + .collect(); + + if commands.is_empty() { + return vec![]; + } + + // Extract the initial query text (everything after the leading sentinel or + // `!`). + let initial_query = line + .strip_prefix('/') + .or_else(|| line.strip_prefix(':')) + .or_else(|| line.strip_prefix('!')) + .unwrap_or(line); + + let mut builder = ForgeWidget::select("Command", commands); + if !initial_query.is_empty() { + builder = builder.with_initial_text(initial_query); + } + + match builder.prompt() { + Ok(Some(row)) => { + vec![InputSuggestion { + value: row.0.name, + span: Span::new(0, line.len()), + append_whitespace: true, + }] + } + _ => vec![], + } + } +} diff --git a/crates/forge_main/src/completer/input_completer.rs b/crates/forge_main/src/completer/input_completer.rs new file mode 100644 index 0000000000000000000000000000000000000000..d1df03d84b9142a81f4049eeacc5e8cc7d3afa07 --- /dev/null +++ b/crates/forge_main/src/completer/input_completer.rs @@ -0,0 +1,102 @@ +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use forge_select::{ForgeWidget, PreviewLayout, PreviewPlacement, SelectRow}; +use forge_walker::Walker; + +use crate::completer::CommandCompleter; +use crate::completer::search_term::{SearchTerm, Span}; +use crate::model::ForgeCommandManager; + +pub fn select_workspace_file(cwd: &Path, query: Option) -> anyhow::Result> { + let files: Vec = Walker::max_all() + .cwd(cwd.to_path_buf()) + .get_blocking() + .unwrap_or_default() + .into_iter() + .map(|file| file.path) + .collect(); + + if files.is_empty() { + return Ok(None); + } + + let has_bat = std::process::Command::new("bat") + .arg("--version") + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .is_ok(); + let cat_cmd = if has_bat { + "bat --color=always --style=numbers,changes --line-range=:500" + } else { + "cat" + }; + + let preview_cmd = format!( + "if [ -d {{}} ]; then ls -la --color=always {{}} 2>/dev/null || ls -la {{}}; else {cat_cmd} {{}}; fi" + ); + let rows: Vec = files + .into_iter() + .map(|path| SelectRow { + raw: path.clone(), + display: path.clone(), + search: path.clone(), + fields: vec![path], + }) + .collect(); + + Ok(ForgeWidget::select_rows("File ❯ ", rows) + .query(Some(query.unwrap_or_default())) + .preview(Some(preview_cmd)) + .preview_layout(PreviewLayout { placement: PreviewPlacement::Bottom, percent: 75 }) + .prompt()? + .map(|row| row.raw)) +} + +pub struct InputCompleter { + cwd: PathBuf, + command: CommandCompleter, +} + +pub struct InputSuggestion { + pub value: String, + pub span: Span, + pub append_whitespace: bool, +} + +impl InputCompleter { + pub fn new(cwd: PathBuf, command_manager: Arc) -> Self { + Self { cwd, command: CommandCompleter::new(command_manager) } + } + + pub fn complete(&mut self, line: &str, pos: usize) -> Vec { + if line.starts_with('/') || line.starts_with(':') { + // if the line starts with '/' or ':' it's probably a command, so we delegate to + // the command completer. + let result = self.command.complete(line, pos); + if !result.is_empty() { + return result; + } + } + + if let Some(query) = SearchTerm::new(line, pos).process() { + let initial_text = if !query.term.is_empty() { + Some(query.term.to_string()) + } else { + None + }; + + if let Ok(Some(selected)) = select_workspace_file(&self.cwd, initial_text) { + let value = format!("[{}]", selected); + return vec![InputSuggestion { + value, + span: Span::new(query.span.start, query.span.end), + append_whitespace: true, + }]; + } + } + + vec![] + } +} diff --git a/crates/forge_main/src/completer/mod.rs b/crates/forge_main/src/completer/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..a650a27349efe8d96015e94d19778fa5de78879c --- /dev/null +++ b/crates/forge_main/src/completer/mod.rs @@ -0,0 +1,6 @@ +mod command; +mod input_completer; +mod search_term; + +pub use command::CommandCompleter; +pub use input_completer::{InputCompleter, select_workspace_file}; diff --git a/crates/forge_main/src/completer/search_term.rs b/crates/forge_main/src/completer/search_term.rs new file mode 100644 index 0000000000000000000000000000000000000000..f3e38f79801b714843d54bc57feb9d33d58e937e --- /dev/null +++ b/crates/forge_main/src/completer/search_term.rs @@ -0,0 +1,128 @@ +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Span { + pub start: usize, + pub end: usize, +} + +impl Span { + pub fn new(start: usize, end: usize) -> Self { + Self { start, end } + } +} + +pub struct SearchTerm { + line: String, + position: usize, +} + +impl SearchTerm { + pub fn new(line: &str, position: usize) -> Self { + if position > line.len() { + panic!( + "Position {position} is out of bounds: string '{line}' (length: {})", + line.len() + ); + } + Self { line: line.to_string(), position } + } + + /// Get the search term from the line based on '@' marker or cursor position + /// + /// If '@' marker is present, returns the word following it. + /// Otherwise, returns the word at the cursor position. + /// If no word is found, returns None. + pub fn process(&self) -> Option> { + // Ensure position is on a UTF-8 character boundary to prevent panics + let safe_position = if self.line.is_char_boundary(self.position) { + self.position + } else { + // Find the nearest lower character boundary + (0..self.position) + .rev() + .find(|&i| self.line.is_char_boundary(i)) + .unwrap_or(0) + }; + + let prefix = self.line.get(..safe_position)?; + let at_pos = prefix.rfind('@')?; + let start_pos = at_pos + 1; + let term = self.line.get(start_pos..safe_position)?; + + Some(TermResult { span: Span::new(start_pos, safe_position), term }) + } +} + +#[derive(Debug)] +pub struct TermResult<'a> { + pub span: Span, + pub term: &'a str, +} + +#[cfg(test)] +mod tests { + use insta::assert_debug_snapshot; + + use super::SearchTerm; + + impl SearchTerm { + fn test(line: &str) -> Vec { + // Test at each valid character boundary position, starting from 1 + (1..=line.len()) + .filter(|&pos| line.is_char_boundary(pos)) + .map(|pos| { + let input = SearchTerm::new(line, pos); + let output = input.process(); + let (a, b) = line.split_at(pos); + + TermSpec { + pos, + input: format!("{a}[{b}"), + output: output.as_ref().map(|term| term.term.to_string()), + span_start: output.as_ref().map(|term| term.span.start), + span_end: output.as_ref().map(|term| term.span.end), + } + }) + .collect() + } + } + + #[derive(Debug)] + #[allow(dead_code)] // Used to generate test snapshots + struct TermSpec { + input: String, + output: Option, + span_start: Option, + span_end: Option, + pos: usize, + } + + #[test] + fn test_marker_based_search() { + let results = SearchTerm::test("@abc @def ghi@"); + assert_debug_snapshot!(results); + } + + #[test] + fn test_marker_based_search_chinese() { + let results = SearchTerm::test("@你好 @世界 测试@"); + assert_debug_snapshot!(results); + } + + #[test] + fn test_marker_based_search_mixed_chinese_english() { + let results = SearchTerm::test("@hello @世界 test@中文"); + assert_debug_snapshot!(results); + } + + #[test] + fn test_marker_based_search_chinese_with_spaces() { + let results = SearchTerm::test("@中 文 @测试"); + assert_debug_snapshot!(results); + } + + #[test] + fn test_marker_based_search_emoji() { + let results = SearchTerm::test("@🚀 @🌟 emoji@"); + assert_debug_snapshot!(results); + } +} diff --git a/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search.snap b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search.snap new file mode 100644 index 0000000000000000000000000000000000000000..711aaa00b7a63385cfde65edcf15ff9930f8553e --- /dev/null +++ b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search.snap @@ -0,0 +1,188 @@ +--- +source: crates/forge_main/src/completer/search_term.rs +expression: results +--- +[ + TermSpec { + input: "@[abc @def ghi@", + output: Some( + "", + ), + span_start: Some( + 1, + ), + span_end: Some( + 1, + ), + pos: 1, + }, + TermSpec { + input: "@a[bc @def ghi@", + output: Some( + "a", + ), + span_start: Some( + 1, + ), + span_end: Some( + 2, + ), + pos: 2, + }, + TermSpec { + input: "@ab[c @def ghi@", + output: Some( + "ab", + ), + span_start: Some( + 1, + ), + span_end: Some( + 3, + ), + pos: 3, + }, + TermSpec { + input: "@abc[ @def ghi@", + output: Some( + "abc", + ), + span_start: Some( + 1, + ), + span_end: Some( + 4, + ), + pos: 4, + }, + TermSpec { + input: "@abc [@def ghi@", + output: Some( + "abc ", + ), + span_start: Some( + 1, + ), + span_end: Some( + 5, + ), + pos: 5, + }, + TermSpec { + input: "@abc @[def ghi@", + output: Some( + "", + ), + span_start: Some( + 6, + ), + span_end: Some( + 6, + ), + pos: 6, + }, + TermSpec { + input: "@abc @d[ef ghi@", + output: Some( + "d", + ), + span_start: Some( + 6, + ), + span_end: Some( + 7, + ), + pos: 7, + }, + TermSpec { + input: "@abc @de[f ghi@", + output: Some( + "de", + ), + span_start: Some( + 6, + ), + span_end: Some( + 8, + ), + pos: 8, + }, + TermSpec { + input: "@abc @def[ ghi@", + output: Some( + "def", + ), + span_start: Some( + 6, + ), + span_end: Some( + 9, + ), + pos: 9, + }, + TermSpec { + input: "@abc @def [ghi@", + output: Some( + "def ", + ), + span_start: Some( + 6, + ), + span_end: Some( + 10, + ), + pos: 10, + }, + TermSpec { + input: "@abc @def g[hi@", + output: Some( + "def g", + ), + span_start: Some( + 6, + ), + span_end: Some( + 11, + ), + pos: 11, + }, + TermSpec { + input: "@abc @def gh[i@", + output: Some( + "def gh", + ), + span_start: Some( + 6, + ), + span_end: Some( + 12, + ), + pos: 12, + }, + TermSpec { + input: "@abc @def ghi[@", + output: Some( + "def ghi", + ), + span_start: Some( + 6, + ), + span_end: Some( + 13, + ), + pos: 13, + }, + TermSpec { + input: "@abc @def ghi@[", + output: Some( + "", + ), + span_start: Some( + 14, + ), + span_end: Some( + 14, + ), + pos: 14, + }, +] diff --git a/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_chinese.snap b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_chinese.snap new file mode 100644 index 0000000000000000000000000000000000000000..e5722f657d74fc4aaf2d3de18528ab2bb2124611 --- /dev/null +++ b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_chinese.snap @@ -0,0 +1,149 @@ +--- +source: crates/forge_main/src/completer/search_term.rs +expression: results +--- +[ + TermSpec { + input: "@[你好 @世界 测试@", + output: Some( + "", + ), + span_start: Some( + 1, + ), + span_end: Some( + 1, + ), + pos: 1, + }, + TermSpec { + input: "@你[好 @世界 测试@", + output: Some( + "你", + ), + span_start: Some( + 1, + ), + span_end: Some( + 4, + ), + pos: 4, + }, + TermSpec { + input: "@你好[ @世界 测试@", + output: Some( + "你好", + ), + span_start: Some( + 1, + ), + span_end: Some( + 7, + ), + pos: 7, + }, + TermSpec { + input: "@你好 [@世界 测试@", + output: Some( + "你好 ", + ), + span_start: Some( + 1, + ), + span_end: Some( + 8, + ), + pos: 8, + }, + TermSpec { + input: "@你好 @[世界 测试@", + output: Some( + "", + ), + span_start: Some( + 9, + ), + span_end: Some( + 9, + ), + pos: 9, + }, + TermSpec { + input: "@你好 @世[界 测试@", + output: Some( + "世", + ), + span_start: Some( + 9, + ), + span_end: Some( + 12, + ), + pos: 12, + }, + TermSpec { + input: "@你好 @世界[ 测试@", + output: Some( + "世界", + ), + span_start: Some( + 9, + ), + span_end: Some( + 15, + ), + pos: 15, + }, + TermSpec { + input: "@你好 @世界 [测试@", + output: Some( + "世界 ", + ), + span_start: Some( + 9, + ), + span_end: Some( + 16, + ), + pos: 16, + }, + TermSpec { + input: "@你好 @世界 测[试@", + output: Some( + "世界 测", + ), + span_start: Some( + 9, + ), + span_end: Some( + 19, + ), + pos: 19, + }, + TermSpec { + input: "@你好 @世界 测试[@", + output: Some( + "世界 测试", + ), + span_start: Some( + 9, + ), + span_end: Some( + 22, + ), + pos: 22, + }, + TermSpec { + input: "@你好 @世界 测试@[", + output: Some( + "", + ), + span_start: Some( + 23, + ), + span_end: Some( + 23, + ), + pos: 23, + }, +] diff --git a/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_chinese_with_spaces.snap b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_chinese_with_spaces.snap new file mode 100644 index 0000000000000000000000000000000000000000..5f8cdd50bad633a2471b08a89ee5cc9889cbd888 --- /dev/null +++ b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_chinese_with_spaces.snap @@ -0,0 +1,110 @@ +--- +source: crates/forge_main/src/completer/search_term.rs +expression: results +--- +[ + TermSpec { + input: "@[中 文 @测试", + output: Some( + "", + ), + span_start: Some( + 1, + ), + span_end: Some( + 1, + ), + pos: 1, + }, + TermSpec { + input: "@中[ 文 @测试", + output: Some( + "中", + ), + span_start: Some( + 1, + ), + span_end: Some( + 4, + ), + pos: 4, + }, + TermSpec { + input: "@中 [文 @测试", + output: Some( + "中 ", + ), + span_start: Some( + 1, + ), + span_end: Some( + 5, + ), + pos: 5, + }, + TermSpec { + input: "@中 文[ @测试", + output: Some( + "中 文", + ), + span_start: Some( + 1, + ), + span_end: Some( + 8, + ), + pos: 8, + }, + TermSpec { + input: "@中 文 [@测试", + output: Some( + "中 文 ", + ), + span_start: Some( + 1, + ), + span_end: Some( + 9, + ), + pos: 9, + }, + TermSpec { + input: "@中 文 @[测试", + output: Some( + "", + ), + span_start: Some( + 10, + ), + span_end: Some( + 10, + ), + pos: 10, + }, + TermSpec { + input: "@中 文 @测[试", + output: Some( + "测", + ), + span_start: Some( + 10, + ), + span_end: Some( + 13, + ), + pos: 13, + }, + TermSpec { + input: "@中 文 @测试[", + output: Some( + "测试", + ), + span_start: Some( + 10, + ), + span_end: Some( + 16, + ), + pos: 16, + }, +] diff --git a/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_emoji.snap b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_emoji.snap new file mode 100644 index 0000000000000000000000000000000000000000..d948576d7dce95ceb77ad7db464223973e8b8e9f --- /dev/null +++ b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_emoji.snap @@ -0,0 +1,162 @@ +--- +source: crates/forge_main/src/completer/search_term.rs +expression: results +--- +[ + TermSpec { + input: "@[🚀 @🌟 emoji@", + output: Some( + "", + ), + span_start: Some( + 1, + ), + span_end: Some( + 1, + ), + pos: 1, + }, + TermSpec { + input: "@🚀[ @🌟 emoji@", + output: Some( + "🚀", + ), + span_start: Some( + 1, + ), + span_end: Some( + 5, + ), + pos: 5, + }, + TermSpec { + input: "@🚀 [@🌟 emoji@", + output: Some( + "🚀 ", + ), + span_start: Some( + 1, + ), + span_end: Some( + 6, + ), + pos: 6, + }, + TermSpec { + input: "@🚀 @[🌟 emoji@", + output: Some( + "", + ), + span_start: Some( + 7, + ), + span_end: Some( + 7, + ), + pos: 7, + }, + TermSpec { + input: "@🚀 @🌟[ emoji@", + output: Some( + "🌟", + ), + span_start: Some( + 7, + ), + span_end: Some( + 11, + ), + pos: 11, + }, + TermSpec { + input: "@🚀 @🌟 [emoji@", + output: Some( + "🌟 ", + ), + span_start: Some( + 7, + ), + span_end: Some( + 12, + ), + pos: 12, + }, + TermSpec { + input: "@🚀 @🌟 e[moji@", + output: Some( + "🌟 e", + ), + span_start: Some( + 7, + ), + span_end: Some( + 13, + ), + pos: 13, + }, + TermSpec { + input: "@🚀 @🌟 em[oji@", + output: Some( + "🌟 em", + ), + span_start: Some( + 7, + ), + span_end: Some( + 14, + ), + pos: 14, + }, + TermSpec { + input: "@🚀 @🌟 emo[ji@", + output: Some( + "🌟 emo", + ), + span_start: Some( + 7, + ), + span_end: Some( + 15, + ), + pos: 15, + }, + TermSpec { + input: "@🚀 @🌟 emoj[i@", + output: Some( + "🌟 emoj", + ), + span_start: Some( + 7, + ), + span_end: Some( + 16, + ), + pos: 16, + }, + TermSpec { + input: "@🚀 @🌟 emoji[@", + output: Some( + "🌟 emoji", + ), + span_start: Some( + 7, + ), + span_end: Some( + 17, + ), + pos: 17, + }, + TermSpec { + input: "@🚀 @🌟 emoji@[", + output: Some( + "", + ), + span_start: Some( + 18, + ), + span_end: Some( + 18, + ), + pos: 18, + }, +] diff --git a/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_mixed_chinese_english.snap b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_mixed_chinese_english.snap new file mode 100644 index 0000000000000000000000000000000000000000..32d1aec5c7bd793109d2ce4d10be67a0af685f1f --- /dev/null +++ b/crates/forge_main/src/completer/snapshots/forge_main__completer__search_term__tests__marker_based_search_mixed_chinese_english.snap @@ -0,0 +1,240 @@ +--- +source: crates/forge_main/src/completer/search_term.rs +expression: results +--- +[ + TermSpec { + input: "@[hello @世界 test@中文", + output: Some( + "", + ), + span_start: Some( + 1, + ), + span_end: Some( + 1, + ), + pos: 1, + }, + TermSpec { + input: "@h[ello @世界 test@中文", + output: Some( + "h", + ), + span_start: Some( + 1, + ), + span_end: Some( + 2, + ), + pos: 2, + }, + TermSpec { + input: "@he[llo @世界 test@中文", + output: Some( + "he", + ), + span_start: Some( + 1, + ), + span_end: Some( + 3, + ), + pos: 3, + }, + TermSpec { + input: "@hel[lo @世界 test@中文", + output: Some( + "hel", + ), + span_start: Some( + 1, + ), + span_end: Some( + 4, + ), + pos: 4, + }, + TermSpec { + input: "@hell[o @世界 test@中文", + output: Some( + "hell", + ), + span_start: Some( + 1, + ), + span_end: Some( + 5, + ), + pos: 5, + }, + TermSpec { + input: "@hello[ @世界 test@中文", + output: Some( + "hello", + ), + span_start: Some( + 1, + ), + span_end: Some( + 6, + ), + pos: 6, + }, + TermSpec { + input: "@hello [@世界 test@中文", + output: Some( + "hello ", + ), + span_start: Some( + 1, + ), + span_end: Some( + 7, + ), + pos: 7, + }, + TermSpec { + input: "@hello @[世界 test@中文", + output: Some( + "", + ), + span_start: Some( + 8, + ), + span_end: Some( + 8, + ), + pos: 8, + }, + TermSpec { + input: "@hello @世[界 test@中文", + output: Some( + "世", + ), + span_start: Some( + 8, + ), + span_end: Some( + 11, + ), + pos: 11, + }, + TermSpec { + input: "@hello @世界[ test@中文", + output: Some( + "世界", + ), + span_start: Some( + 8, + ), + span_end: Some( + 14, + ), + pos: 14, + }, + TermSpec { + input: "@hello @世界 [test@中文", + output: Some( + "世界 ", + ), + span_start: Some( + 8, + ), + span_end: Some( + 15, + ), + pos: 15, + }, + TermSpec { + input: "@hello @世界 t[est@中文", + output: Some( + "世界 t", + ), + span_start: Some( + 8, + ), + span_end: Some( + 16, + ), + pos: 16, + }, + TermSpec { + input: "@hello @世界 te[st@中文", + output: Some( + "世界 te", + ), + span_start: Some( + 8, + ), + span_end: Some( + 17, + ), + pos: 17, + }, + TermSpec { + input: "@hello @世界 tes[t@中文", + output: Some( + "世界 tes", + ), + span_start: Some( + 8, + ), + span_end: Some( + 18, + ), + pos: 18, + }, + TermSpec { + input: "@hello @世界 test[@中文", + output: Some( + "世界 test", + ), + span_start: Some( + 8, + ), + span_end: Some( + 19, + ), + pos: 19, + }, + TermSpec { + input: "@hello @世界 test@[中文", + output: Some( + "", + ), + span_start: Some( + 20, + ), + span_end: Some( + 20, + ), + pos: 20, + }, + TermSpec { + input: "@hello @世界 test@中[文", + output: Some( + "中", + ), + span_start: Some( + 20, + ), + span_end: Some( + 23, + ), + pos: 23, + }, + TermSpec { + input: "@hello @世界 test@中文[", + output: Some( + "中文", + ), + span_start: Some( + 20, + ), + span_end: Some( + 26, + ), + pos: 26, + }, +] diff --git a/crates/forge_main/src/conversation_selector.rs b/crates/forge_main/src/conversation_selector.rs new file mode 100644 index 0000000000000000000000000000000000000000..ea755621e68f2319c26f16dd035d4c044d748e86 --- /dev/null +++ b/crates/forge_main/src/conversation_selector.rs @@ -0,0 +1,186 @@ +use anyhow::Result; +use chrono::Utc; +use forge_api::Conversation; +use forge_domain::ConversationId; +use forge_select::{ForgeWidget, PreviewLayout, PreviewPlacement, SelectRow}; + +use crate::display_constants::markers; +use crate::info::Info; +use crate::porcelain::Porcelain; + +/// Logic for selecting conversations from a list +pub struct ConversationSelector; + +impl ConversationSelector { + /// Select a conversation from the provided list using a custom TUI with + /// a preview pane showing conversation details. + /// + /// The preview command uses `forge conversation info` and + /// `forge conversation show` to display the selected conversation's + /// metadata and last message side-by-side with the picker list. + /// + /// Returns the selected conversation, or None if the user cancelled. + pub async fn select_conversation( + conversations: &[Conversation], + _current_conversation_id: Option, + query: Option, + ) -> Result> { + if conversations.is_empty() { + return Ok(None); + } + + // Filter to conversations with titles and context + let valid_conversations: Vec<&Conversation> = conversations + .iter() + .filter(|c| c.context.is_some()) + .collect(); + + if valid_conversations.is_empty() { + return Ok(None); + } + + // Build Info structure for display + let now = Utc::now(); + let mut info = Info::new(); + + for conv in &valid_conversations { + let title = conv + .title + .as_deref() + .map(|t| t.to_string()) + .unwrap_or_else(|| markers::EMPTY.to_string()); + + let duration = now.signed_duration_since( + conv.metadata.updated_at.unwrap_or(conv.metadata.created_at), + ); + let duration = + std::time::Duration::from_secs((duration.num_minutes() * 60).max(0) as u64); + let time_ago = if duration.is_zero() { + "now".to_string() + } else { + format!("{} ago", humantime::format_duration(duration)) + }; + + info = info + .add_title(conv.id) + .add_key_value("Title", title) + .add_key_value("Updated", time_ago); + } + + // Convert to porcelain, drop the UUID title column (col 0), truncate the + // Title column for display, uppercase headers + let porcelain_output = Porcelain::from(&info) + .drop_col(0) + .truncate(0, 60) + .uppercase_headers(); + let porcelain_str = porcelain_output.to_string(); + + let all_lines: Vec<&str> = porcelain_str.lines().collect(); + if all_lines.is_empty() { + return Ok(None); + } + + // Build SelectRow items for the shared Rust selector UI. + // Each row stores the UUID in `fields[0]` so that `{1}` in the preview + // command resolves to the conversation ID. The `raw` field is what gets + // returned on selection (the UUID). + let mut rows: Vec = Vec::with_capacity(all_lines.len()); + + // Header row (non-selectable via header_lines=1) + if let Some(header) = all_lines.first() { + rows.push(SelectRow::header(header.to_string())); + } + + // Data rows: each maps to a conversation + for (i, line) in all_lines.iter().skip(1).enumerate() { + if let Some(conv) = valid_conversations.get(i) { + let uuid = conv.id.to_string(); + rows.push(SelectRow { + raw: uuid.clone(), + display: line.to_string(), + search: line.to_string(), + fields: vec![uuid], + }); + } + } + + // Build a lookup map from UUID to Conversation for the result + let conv_map: std::collections::HashMap = valid_conversations + .into_iter() + .map(|c| (c.id.to_string(), c.clone())) + .collect(); + + let preview_command = + "CLICOLOR_FORCE=1 forge conversation info {1}; echo; CLICOLOR_FORCE=1 forge conversation show {1}" + .to_string(); + + let selected_uuid = tokio::task::spawn_blocking(move || -> Result> { + Ok(ForgeWidget::select_rows("Conversation", rows) + .query(query) + .header_lines(1_usize) + .preview(Some(preview_command)) + .preview_layout(PreviewLayout { placement: PreviewPlacement::Bottom, percent: 60 }) + .prompt()? + .map(|row| row.raw)) + }) + .await??; + + Ok(selected_uuid.and_then(|uuid| conv_map.get(&uuid).cloned())) + } +} + +#[cfg(test)] +mod tests { + use chrono::Utc; + use forge_api::Conversation; + use forge_domain::{ConversationId, MetaData, Metrics}; + use pretty_assertions::assert_eq; + + use super::*; + + fn create_test_conversation(id: &str, title: Option<&str>) -> Conversation { + let now = Utc::now(); + Conversation { + id: ConversationId::parse(id).unwrap(), + title: title.map(|t| t.to_string()), + context: None, + metrics: Metrics::default().started_at(now), + metadata: MetaData { created_at: now, updated_at: Some(now) }, + } + } + + #[tokio::test] + async fn test_select_conversation_empty_list() { + let conversations = vec![]; + let result = ConversationSelector::select_conversation(&conversations, None, None) + .await + .unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_select_conversation_with_titles() { + let conversations = [ + create_test_conversation( + "550e8400-e29b-41d4-a716-446655440000", + Some("First Conversation"), + ), + create_test_conversation( + "550e8400-e29b-41d4-a716-446655440001", + Some("Second Conversation"), + ), + ]; + + assert_eq!(conversations.len(), 2); + } + + #[test] + fn test_select_conversation_without_titles() { + let conversations = [ + create_test_conversation("550e8400-e29b-41d4-a716-446655440002", None), + create_test_conversation("550e8400-e29b-41d4-a716-446655440003", None), + ]; + + assert_eq!(conversations.len(), 2); + } +} diff --git a/crates/forge_main/src/display_constants.rs b/crates/forge_main/src/display_constants.rs new file mode 100644 index 0000000000000000000000000000000000000000..230bfb7578f3e899766a2efd1715f6428b14949e --- /dev/null +++ b/crates/forge_main/src/display_constants.rs @@ -0,0 +1,143 @@ +//! Centralized display constants for consistent UI messaging +//! +//! This module provides a single source of truth for all display values used +//! across the application. When adding new constants, follow these conventions: +//! +//! - **Placeholders** (missing/unknown data): Use angle brackets `` +//! - **Special markers**: Use square brackets `[value]` +//! - **Status values**: Use lowercase for user-facing strings +//! - **Type discriminators**: Use enums instead of string constants +//! +//! # Examples +//! +//! ```rust,ignore +//! use crate::display_constants::{placeholders, status}; +//! +//! // For special values +//! let empty = placeholders::EMPTY; // "[empty]" +//! +//! // For status indicators +//! info.add_key_value("status", status::ENABLED); // "[enabled]" +//! ``` + +use std::fmt; + +/// Status indicator values. +/// +/// Use lowercase for user-facing status strings to maintain consistency. +pub mod status { + /// Indicates a resource is enabled/configured + pub const YES: &str = "[yes]"; + + /// Indicates a resource is disabled + pub const NO: &str = "[no]"; +} + +/// Table column headers for porcelain (machine-readable) output. +/// +/// These headers use the `$` prefix to distinguish them as metadata columns. +pub mod headers { + /// Default ID column header + pub const ID: &str = "ID"; + + /// Field name column header + pub const FIELD: &str = "FIELD"; + + /// Field value column header + pub const VALUE: &str = "VALUE"; +} + +/// Special markers for specific contexts. +/// +/// These use square brackets to indicate special/synthetic values. +pub mod markers { + /// Indicates an empty value (distinct from null/unset) + pub const EMPTY: &str = "[empty]"; + + /// Indicates a built-in (non-user-defined) component + pub const BUILT_IN: &str = "[built-in]"; +} + +/// Type discriminator for commands, agents, and custom entries. +/// +/// Use this enum instead of string constants for type-safe discrimination +/// between command types in listings and displays. +/// +/// # Examples +/// +/// ```rust,ignore +/// let cmd_type = CommandType::Agent; +/// info.add_key_value("type", cmd_type.as_str()); +/// ``` +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CommandType { + /// Built-in command + Command, + /// Agent (AI assistant with specific role) + Agent, + /// User-defined custom command + Custom, +} + +impl CommandType { + /// Returns the string representation of the command type. + /// + /// This is the canonical way to convert a CommandType to a string + /// for display purposes. + pub const fn as_str(&self) -> &'static str { + match self { + Self::Command => "command", + Self::Agent => "agent", + Self::Custom => "custom", + } + } +} + +impl fmt::Display for CommandType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_command_type_display() { + assert_eq!(CommandType::Command.to_string(), "command"); + assert_eq!(CommandType::Agent.to_string(), "agent"); + assert_eq!(CommandType::Custom.to_string(), "custom"); + } + + #[test] + fn test_command_type_as_str() { + assert_eq!(CommandType::Command.as_str(), "command"); + assert_eq!(CommandType::Agent.as_str(), "agent"); + assert_eq!(CommandType::Custom.as_str(), "custom"); + } + + #[test] + fn test_placeholders_use_square_brackets() { + // EMPTY uses square brackets like other special markers + assert!(markers::EMPTY.starts_with('[')); + assert!(markers::EMPTY.ends_with(']')); + } + + #[test] + fn test_markers_have_square_brackets() { + assert!(markers::BUILT_IN.starts_with('[')); + assert!(markers::BUILT_IN.ends_with(']')); + assert!(markers::EMPTY.starts_with('[')); + assert!(markers::EMPTY.ends_with(']')); + } + + #[test] + fn test_status_values_use_square_brackets() { + // Status values use square brackets to distinguish them from raw strings + assert!(status::YES.starts_with('[')); + assert!(status::YES.ends_with(']')); + assert!(status::NO.starts_with('[')); + assert!(status::NO.ends_with(']')); + } +} diff --git a/crates/forge_main/src/editor.rs b/crates/forge_main/src/editor.rs new file mode 100644 index 0000000000000000000000000000000000000000..0f96a614922ab2fb0b7c4efe1d1bec9479da7c17 --- /dev/null +++ b/crates/forge_main/src/editor.rs @@ -0,0 +1,309 @@ +use std::borrow::Cow; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +use console::{measure_text_width, strip_ansi_codes}; +use forge_api::Environment; +use nu_ansi_term::Style; +use rustyline::completion::{Completer, Pair}; +use rustyline::config::{ColorMode, CompletionType, Config}; +use rustyline::error::ReadlineError as RustyReadlineError; +use rustyline::highlight::Highlighter; +use rustyline::hint::{Hinter, HistoryHinter}; +use rustyline::history::DefaultHistory; +use rustyline::validate::Validator; +use rustyline::{ + Cmd, Context as RustylineContext, Editor, EventHandler, Helper, KeyCode, KeyEvent, Modifiers, + Prompt as RustylinePrompt, +}; + +use super::completer::InputCompleter; +use super::zsh::paste::wrap_pasted_text; +use crate::highlighter::ForgeHighlighter; +use crate::model::ForgeCommandManager; +use crate::prompt::ForgePrompt; + +const HISTORY_CAPACITY: usize = 1024 * 1024; + +/// Interactive terminal editor used by the Forge prompt. +pub struct ForgeEditor { + editor: Editor, + history_file: PathBuf, + pending_buffer: Option, +} + +/// Result of reading one prompt interaction from the terminal. +#[derive(Debug, PartialEq, Eq)] +pub enum ReadResult { + Success(String), + Empty, + Continue, + Exit, +} + +impl ForgeEditor { + /// Creates a new interactive editor with history, completion, and + /// highlighting. + pub fn new( + env: Environment, + custom_history_path: Option, + manager: Arc, + ) -> Self { + let history_file = env.history_path(custom_history_path.as_ref()); + let helper = ForgeHelper::new(env.cwd, manager); + let config = Config::builder() + .max_history_size(HISTORY_CAPACITY) + .expect("rustyline history capacity should be valid") + .completion_type(CompletionType::List) + .completion_show_all_if_ambiguous(true) + .color_mode(ColorMode::Forced) + .enable_signals(true) + .build(); + let mut editor = Editor::::with_config(config) + .expect("rustyline editor should initialize for an interactive terminal"); + editor.bind_sequence( + KeyEvent(KeyCode::Enter, Modifiers::ALT), + EventHandler::Simple(Cmd::Newline), + ); + editor.bind_sequence( + KeyEvent(KeyCode::Char('k'), Modifiers::CTRL), + EventHandler::Simple(Cmd::ClearScreen), + ); + editor.bind_sequence( + KeyEvent(KeyCode::Char('K'), Modifiers::CTRL), + EventHandler::Simple(Cmd::ClearScreen), + ); + editor.set_helper(Some(helper)); + let _ = editor.load_history(&history_file); + Self { editor, history_file, pending_buffer: None } + } + + fn normalize_result(&mut self, buffer: String) -> ReadResult { + let result = normalize_result_text(buffer); + if let ReadResult::Success(text) = &result { + let _ = self.editor.add_history_entry(text.as_str()); + let _ = self.editor.save_history(&self.history_file); + } + result + } + + /// Reads one logical input from the terminal. + pub fn prompt(&mut self, prompt: &mut ForgePrompt) -> anyhow::Result { + let prompt_text = render_prompt(prompt); + let initial = self.pending_buffer.take().unwrap_or_default(); + let readline = if initial.is_empty() { + self.editor.readline(&prompt_text) + } else { + self.editor + .readline_with_initial(&prompt_text, (&initial, "")) + }; + prompt.refresh(); + + match readline { + Ok(buffer) => Ok(self.normalize_result(buffer)), + Err(RustyReadlineError::Interrupted) => Ok(ReadResult::Continue), + Err(RustyReadlineError::Eof) => Ok(ReadResult::Exit), + Err(error) => Err(anyhow::anyhow!(ReadLineError(error))), + } + } + + /// Sets the buffer content to be pre-filled on the next prompt. + pub fn set_buffer(&mut self, content: String) { + self.pending_buffer = Some(content); + } +} + +#[derive(Debug, thiserror::Error)] +#[error("failed to read line from terminal: {0}")] +pub struct ReadLineError(RustyReadlineError); + +fn normalize_result_text(buffer: String) -> ReadResult { + let trimmed = buffer.trim(); + if trimmed.is_empty() { + return ReadResult::Empty; + } + ReadResult::Success(wrap_pasted_text(trimmed)) +} + +fn render_prompt(prompt: &ForgePrompt) -> ResponsivePrompt { + let left = prompt.render_prompt_left(); + let indicator = prompt.render_prompt_indicator(); + let right = prompt.render_prompt_right(); + let right = right.trim_start(); + + // `raw` is what rustyline measures to position the cursor; `styled` is what + // it prints. `raw` MUST be free of ANSI escapes: rustyline's Windows + // console backend computes cursor columns by counting grapheme widths of + // `raw` (it cannot interpret escape sequences and debug-asserts against + // them), so any styling left in `raw` is counted as visible width and + // pushes the cursor past where the text actually is. The left prompt and + // indicator are styled via `nu_ansi_term`, so strip those codes for `raw`. + // The right prompt is positioned off to the side with cursor save/restore + // and is not part of the input-line geometry, so it is excluded from `raw` + // entirely. + if right.trim().is_empty() { + let styled = format!("{left}{indicator}"); + let raw = strip_ansi_codes(&styled).into_owned(); + return ResponsivePrompt { raw, styled }; + } + + if let Some((first_line, remaining)) = left.split_once('\n') { + let right = render_right_prompt(right); + let raw = strip_ansi_codes(&format!("{first_line}\n{remaining}{indicator}")).into_owned(); + return ResponsivePrompt { + raw, + styled: format!("{first_line}{right}\n{remaining}{indicator}"), + }; + } + + let right = render_right_prompt(right); + let raw = strip_ansi_codes(&format!("{left}{indicator}")).into_owned(); + ResponsivePrompt { raw, styled: format!("{left}{right}{indicator}") } +} + +fn render_right_prompt(right: &str) -> String { + let width = measure_text_width(strip_ansi_codes(right).as_ref()); + format!("\x1b[s\x1b[999C\x1b[{width}D{right}\x1b[K\x1b[u") +} + +struct ResponsivePrompt { + raw: String, + styled: String, +} + +impl RustylinePrompt for ResponsivePrompt { + fn raw(&self) -> &str { + &self.raw + } + + fn styled(&self) -> &str { + &self.styled + } +} + +struct ForgeHelper { + completer: Mutex, + highlighter: ForgeHighlighter, + hinter: HistoryHinter, +} + +impl ForgeHelper { + fn new(cwd: PathBuf, command_manager: Arc) -> Self { + Self { + completer: Mutex::new(InputCompleter::new(cwd, command_manager)), + highlighter: ForgeHighlighter, + hinter: HistoryHinter {}, + } + } +} + +impl Helper for ForgeHelper {} + +impl Completer for ForgeHelper { + type Candidate = Pair; + + fn complete( + &self, + line: &str, + pos: usize, + _ctx: &RustylineContext<'_>, + ) -> rustyline::Result<(usize, Vec)> { + let mut completer = self + .completer + .lock() + .expect("input completer mutex poisoned"); + let suggestions = completer.complete(line, pos); + let start = suggestions + .iter() + .map(|suggestion| suggestion.span.start) + .min() + .unwrap_or(pos); + let pairs = suggestions + .into_iter() + .map(|suggestion| { + let replacement = if suggestion.append_whitespace { + format!("{} ", suggestion.value) + } else { + suggestion.value + }; + Pair { display: replacement.clone(), replacement } + }) + .collect(); + Ok((start, pairs)) + } +} + +impl Hinter for ForgeHelper { + type Hint = String; + + fn hint(&self, line: &str, pos: usize, ctx: &RustylineContext<'_>) -> Option { + self.hinter.hint(line, pos, ctx) + } +} + +impl Highlighter for ForgeHelper { + fn highlight<'l>(&self, line: &'l str, pos: usize) -> Cow<'l, str> { + let styled = self.highlighter.highlight(line, pos); + if styled.buffer.is_empty() { + return Cow::Borrowed(line); + } + + let default_style = Style::new(); + let mut rendered = String::with_capacity(line.len()); + for (style, text) in styled.buffer { + if style == default_style { + rendered.push_str(&text); + } else { + rendered.push_str(&style.paint(text).to_string()); + } + } + Cow::Owned(rendered) + } + + fn highlight_hint<'h>(&self, hint: &'h str) -> Cow<'h, str> { + Cow::Owned(Style::new().dimmed().paint(hint).to_string()) + } +} + +impl Validator for ForgeHelper {} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_normalize_result_wraps_existing_pasted_path() { + let fixture = "/usr/bin/env".to_string(); + + let actual = normalize_result_text(fixture); + + let expected = ReadResult::Success("@[/usr/bin/env]".to_string()); + assert_eq!(actual, expected); + } + + #[test] + fn test_render_prompt_raw_has_no_ansi_escapes() { + use std::path::PathBuf; + + use forge_api::{AgentId, ModelId}; + + // rustyline measures `raw()` to position the cursor and, on Windows, + // cannot interpret ANSI escapes (it counts their bytes as visible + // columns). `raw()` must therefore be free of escape sequences even + // though the visible prompt is styled. + let mut prompt = ForgePrompt::new(PathBuf::from("project"), AgentId::default()); + prompt.model(ModelId::new("anthropic/claude-opus-4")); + + let rendered = render_prompt(&prompt); + + assert!( + !rendered.raw.contains('\u{1b}'), + "raw prompt must not contain ANSI escape sequences: {:?}", + rendered.raw + ); + // The styled prompt, by contrast, does carry styling for display. + assert!(rendered.styled.contains('\u{1b}')); + } +} diff --git a/crates/forge_main/src/error.rs b/crates/forge_main/src/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..58a336f4665ddd1ecfb462a4cfa3beafb102a993 --- /dev/null +++ b/crates/forge_main/src/error.rs @@ -0,0 +1,28 @@ +use forge_domain::ProviderId; + +/// Errors specific to UI operations +#[derive(Debug, thiserror::Error)] +pub enum UIError { + /// No authentication methods are available for a provider + #[error( + "No authentication methods are configured for provider '{provider}'. \ + Please check your provider configuration." + )] + NoAuthMethodsAvailable { provider: ProviderId }, + + /// User selected an authentication method that could not be found + #[error( + "The selected authentication method is no longer available. \ + Please try again or check your provider configuration." + )] + AuthMethodNotFound, + + /// Display data is missing a header line - occurs when the data source + /// (agents, models, or providers list) produces empty output after + /// formatting + #[error( + "Unable to display the selection list - the data appears to be empty. \ + This can happen if the agents, models, or providers list could not be retrieved" + )] + MissingHeaderLine, +} diff --git a/crates/forge_main/src/highlighter.rs b/crates/forge_main/src/highlighter.rs new file mode 100644 index 0000000000000000000000000000000000000000..10919c3c75168e71d5a95dee2cb20209644eaac0 --- /dev/null +++ b/crates/forge_main/src/highlighter.rs @@ -0,0 +1,318 @@ +use nu_ansi_term::{Color, Style}; + +pub(crate) struct StyledText { + pub(crate) buffer: Vec<(Style, String)>, +} + +impl StyledText { + pub fn new() -> Self { + Self { buffer: Vec::new() } + } + + pub fn push(&mut self, value: (Style, String)) { + self.buffer.push(value); + } +} + +/// Syntax highlighter for the forge readline prompt. +/// +/// Applies visual styles to recognised input patterns as the user types: +/// - Commands (`:foo` or `/foo` for backward compatibility) are rendered in +/// yellow bold. +/// - File mentions (`@[path]`) are rendered in cyan bold. +/// - Shell pass-through commands (`!cmd`) are rendered in magenta. +/// - All other text is rendered in the default terminal style. +pub struct ForgeHighlighter; + +impl ForgeHighlighter { + pub(crate) fn highlight(&self, line: &str, _cursor: usize) -> StyledText { + let mut styled = StyledText::new(); + + if line.is_empty() { + return styled; + } + + // Command: highlight the command token (e.g. `:compact` or `/compact` for + // compat) in yellow bold, then the remainder (arguments) without + // special styling. + if line.starts_with('/') || line.starts_with(':') { + let end = line.find(|c: char| c.is_whitespace()).unwrap_or(line.len()); + styled.push(( + Style::new().bold().fg(Color::Yellow), + line.get(..end).unwrap_or(line).to_string(), + )); + if end < line.len() + && let Some(args) = line.get(end..) + { + highlight_mentions(args, &mut styled); + } + return styled; + } + + // Shell pass-through: `!` rendered in magenta. + if line.starts_with('!') { + styled.push((Style::new().fg(Color::Magenta), line.to_string())); + return styled; + } + + // General message text — scan for `@[...]` file mentions and colour them cyan + // bold. + highlight_mentions(line, &mut styled); + + styled + } +} + +/// Walk through `line` and emit styled segments, colouring every `@[...]` +/// mention (matching the ZSH pattern `@\[[^]]*\]`) in cyan bold and leaving +/// surrounding text unstyled. +/// +/// Mirrors `ZSH_HIGHLIGHT_PATTERNS+=('@\[[^]]#\]' 'fg=cyan,bold')` exactly: +/// - Requires `@[` opener. +/// - Matches zero or more non-`]` characters inside the brackets. +/// - Requires closing `]` — unterminated tags are left unstyled. +fn highlight_mentions(line: &str, styled: &mut StyledText) { + let mut remaining = line; + + while !remaining.is_empty() { + // Find the next `@[` opener. + match remaining.find("@[") { + None => { + // No more mentions — emit the rest as plain text. + styled.push((Style::new(), remaining.to_string())); + break; + } + Some(start) => { + // Emit any plain text before the `@[`. + if start > 0 + && let Some(before) = remaining.get(..start) + { + styled.push((Style::new(), before.to_string())); + } + + // `after_open` starts at `@[`. + let after_open = match remaining.get(start..) { + Some(s) => s, + None => { + styled.push((Style::new(), remaining.to_string())); + break; + } + }; + + // Look for a `]` that is not immediately after `@[`. + // The ZSH pattern `[^]#` matches zero-or-more non-`]` chars, + // so `@[]` (empty brackets) also qualifies. + // We search for `]` starting from position 2 (after `@[`). + match after_open.get(2..).and_then(|s| s.find(']')) { + None => { + // No closing `]` — emit `@[` and the rest as plain text + // to match ZSH behaviour (unterminated tag = no highlight). + styled.push((Style::new(), after_open.to_string())); + break; + } + Some(rel_close) => { + // Absolute position of `]` within `after_open`. + let close = 2 + rel_close; + // Emit `@[...]` in cyan bold (inclusive of both brackets). + let mention = after_open.get(..=close).unwrap_or(after_open); + styled.push((Style::new().bold().fg(Color::Cyan), mention.to_string())); + match after_open.get(close + 1..) { + Some(rest) => remaining = rest, + None => break, + } + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn render(styled: &StyledText) -> String { + styled.buffer.iter().map(|(_, s)| s.as_str()).collect() + } + + fn styles(styled: &StyledText) -> Vec