add context window monitor

Writes the current context window to logs/current_context_window (uses a symlink to a session ID).

This PR was unfortunately generated by a different LLM and did a ton of superficial reformating, it's actually a fairly small and benign change, but I don't want to roll back everything. Hope that's ok.
This commit is contained in:
Jochen
2025-11-27 21:00:02 +11:00
parent 93dc4acf86
commit 52f78653b4
89 changed files with 4040 additions and 2576 deletions
Generated
+2
View File
@@ -1351,6 +1351,8 @@ version = "0.1.0"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"g3-cli", "g3-cli",
"g3-providers",
"serde_json",
"tokio", "tokio",
] ]
+6
View File
@@ -45,3 +45,9 @@ license = "MIT"
g3-cli = { path = "crates/g3-cli" } g3-cli = { path = "crates/g3-cli" }
tokio = { workspace = true } tokio = { workspace = true }
anyhow = { workspace = true } anyhow = { workspace = true }
g3-providers = { path = "crates/g3-providers" }
serde_json = { workspace = true }
[[example]]
name = "verify_message_id"
path = "examples/verify_message_id.rs"
+290 -114
View File
@@ -1,5 +1,5 @@
use anyhow::Result; use anyhow::Result;
use crossterm::style::{Color, SetForegroundColor, ResetColor}; use crossterm::style::{Color, ResetColor, SetForegroundColor};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@@ -18,8 +18,13 @@ fn generate_turn_histogram(turn_metrics: &[TurnMetrics]) -> String {
let mut histogram = String::new(); let mut histogram = String::new();
// Find max values for scaling // Find max values for scaling
let max_tokens = turn_metrics.iter().map(|t| t.tokens_used).max().unwrap_or(1); let max_tokens = turn_metrics
let max_time_ms = turn_metrics.iter() .iter()
.map(|t| t.tokens_used)
.max()
.unwrap_or(1);
let max_time_ms = turn_metrics
.iter()
.map(|t| t.wall_clock_time.as_millis().min(u32::MAX as u128) as u32) .map(|t| t.wall_clock_time.as_millis().min(u32::MAX as u128) as u32)
.max() .max()
.unwrap_or(1); .unwrap_or(1);
@@ -30,8 +35,15 @@ fn generate_turn_histogram(turn_metrics: &[TurnMetrics]) -> String {
const TIME_CHAR: char = '▓'; const TIME_CHAR: char = '▓';
histogram.push_str("\n📊 Per-Turn Performance Histogram:\n"); histogram.push_str("\n📊 Per-Turn Performance Histogram:\n");
histogram.push_str(&format!(" {} = Tokens Used (max: {})\n", TOKEN_CHAR, max_tokens)); histogram.push_str(&format!(
histogram.push_str(&format!(" {} = Wall Clock Time (max: {:.1}s)\n\n", TIME_CHAR, max_time_ms as f64 / 1000.0)); " {} = Tokens Used (max: {})\n",
TOKEN_CHAR, max_tokens
));
histogram.push_str(&format!(
" {} = Wall Clock Time (max: {:.1}s)\n\n",
TIME_CHAR,
max_time_ms as f64 / 1000.0
));
for metrics in turn_metrics { for metrics in turn_metrics {
let turn_time_ms = metrics.wall_clock_time.as_millis().min(u32::MAX as u128) as u32; let turn_time_ms = metrics.wall_clock_time.as_millis().min(u32::MAX as u128) as u32;
@@ -67,19 +79,17 @@ fn generate_turn_histogram(turn_metrics: &[TurnMetrics]) -> String {
// Add turn information // Add turn information
histogram.push_str(&format!( histogram.push_str(&format!(
" Turn {:2}: {:>6} tokens │{:<40}│\n", " Turn {:2}: {:>6} tokens │{:<40}│\n",
metrics.turn_number, metrics.turn_number, metrics.tokens_used, token_bar
metrics.tokens_used,
token_bar
)); ));
histogram.push_str(&format!( histogram.push_str(&format!(
" {:>6} │{:<40}│\n", " {:>6} │{:<40}│\n",
time_str, time_str, time_bar
time_bar
)); ));
// Add separator line between turns (except for last turn) // Add separator line between turns (except for last turn)
if metrics.turn_number != turn_metrics.last().unwrap().turn_number { if metrics.turn_number != turn_metrics.last().unwrap().turn_number {
histogram.push_str(" ────────────┼────────────────────────────────────────┤\n"); histogram
.push_str(" ────────────┼────────────────────────────────────────┤\n");
} }
} }
@@ -90,10 +100,20 @@ fn generate_turn_histogram(turn_metrics: &[TurnMetrics]) -> String {
let avg_time_ms = total_time.as_millis() as f64 / turn_metrics.len() as f64; let avg_time_ms = total_time.as_millis() as f64 / turn_metrics.len() as f64;
histogram.push_str("\n📈 Summary Statistics:\n"); histogram.push_str("\n📈 Summary Statistics:\n");
histogram.push_str(&format!(" • Total Tokens: {} across {} turns\n", total_tokens, turn_metrics.len())); histogram.push_str(&format!(
" • Total Tokens: {} across {} turns\n",
total_tokens,
turn_metrics.len()
));
histogram.push_str(&format!(" • Average Tokens/Turn: {:.1}\n", avg_tokens)); histogram.push_str(&format!(" • Average Tokens/Turn: {:.1}\n", avg_tokens));
histogram.push_str(&format!(" • Total Time: {:.1}s\n", total_time.as_secs_f64())); histogram.push_str(&format!(
histogram.push_str(&format!(" • Average Time/Turn: {:.1}s\n", avg_time_ms / 1000.0)); " • Total Time: {:.1}s\n",
total_time.as_secs_f64()
));
histogram.push_str(&format!(
" • Average Time/Turn: {:.1}s\n",
avg_time_ms / 1000.0
));
histogram histogram
} }
@@ -181,15 +201,15 @@ use g3_config::Config;
use g3_core::{project::Project, ui_writer::UiWriter, Agent, DiscoveryOptions}; use g3_core::{project::Project, ui_writer::UiWriter, Agent, DiscoveryOptions};
use rustyline::error::ReadlineError; use rustyline::error::ReadlineError;
use rustyline::DefaultEditor; use rustyline::DefaultEditor;
use sha2::{Digest, Sha256};
use std::path::Path; use std::path::Path;
use std::path::PathBuf; use std::path::PathBuf;
use sha2::{Digest, Sha256};
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
use tracing::{error, info}; use tracing::{error, info};
use g3_core::error_handling::{classify_error, ErrorType, RecoverableError}; use g3_core::error_handling::{classify_error, ErrorType, RecoverableError};
mod ui_writer_impl;
mod simple_output; mod simple_output;
mod ui_writer_impl;
use simple_output::SimpleOutput; use simple_output::SimpleOutput;
mod machine_ui_writer; mod machine_ui_writer;
use machine_ui_writer::MachineUiWriter; use machine_ui_writer::MachineUiWriter;
@@ -285,7 +305,7 @@ pub struct Cli {
/// Enable fast codebase discovery before first LLM turn /// Enable fast codebase discovery before first LLM turn
#[arg(long, value_name = "PATH")] #[arg(long, value_name = "PATH")]
pub codebase_fast_start: Option<PathBuf> pub codebase_fast_start: Option<PathBuf>,
} }
pub async fn run() -> Result<()> { pub async fn run() -> Result<()> {
@@ -293,9 +313,16 @@ pub async fn run() -> Result<()> {
// Check if flock mode is enabled // Check if flock mode is enabled
if let (Some(project_dir), Some(flock_workspace), Some(num_segments)) = if let (Some(project_dir), Some(flock_workspace), Some(num_segments)) =
(&cli.project, &cli.flock_workspace, cli.segments) { (&cli.project, &cli.flock_workspace, cli.segments)
{
// Run flock mode // Run flock mode
return run_flock_mode(project_dir.clone(), flock_workspace.clone(), num_segments, cli.flock_max_turns).await; return run_flock_mode(
project_dir.clone(),
flock_workspace.clone(),
num_segments,
cli.flock_max_turns,
)
.await;
} }
// Otherwise, continue with normal mode // Otherwise, continue with normal mode
@@ -361,7 +388,10 @@ pub async fn run() -> Result<()> {
let project = if cli.autonomous { let project = if cli.autonomous {
if let Some(requirements_text) = &cli.requirements { if let Some(requirements_text) = &cli.requirements {
// Use requirements text override // Use requirements text override
Project::new_autonomous_with_requirements(workspace_dir.clone(), requirements_text.clone())? Project::new_autonomous_with_requirements(
workspace_dir.clone(),
requirements_text.clone(),
)?
} else { } else {
// Use traditional requirements.md file // Use traditional requirements.md file
Project::new_autonomous(workspace_dir.clone())? Project::new_autonomous(workspace_dir.clone())?
@@ -413,9 +443,7 @@ pub async fn run() -> Result<()> {
// Combine AGENTS.md and README content if both exist // Combine AGENTS.md and README content if both exist
let combined_content = match (agents_content.clone(), readme_content.clone()) { let combined_content = match (agents_content.clone(), readme_content.clone()) {
(Some(agents), Some(readme)) => { (Some(agents), Some(readme)) => Some(format!("{}\n\n{}", agents, readme)),
Some(format!("{}\n\n{}", agents, readme))
}
(Some(agents), None) => Some(agents), (Some(agents), None) => Some(agents),
(None, Some(readme)) => Some(readme), (None, Some(readme)) => Some(readme),
(None, None) => None, (None, None) => None,
@@ -458,7 +486,8 @@ pub async fn run() -> Result<()> {
if use_accumulative { if use_accumulative {
// Run accumulative mode and return early // Run accumulative mode and return early
run_accumulative_mode(workspace_dir.clone(), cli.clone(), combined_content.clone()).await?; run_accumulative_mode(workspace_dir.clone(), cli.clone(), combined_content.clone())
.await?;
return Ok(()); return Ok(());
} }
@@ -534,9 +563,12 @@ async fn run_accumulative_mode(
output.print("g3 programming agent - autonomous mode"); output.print("g3 programming agent - autonomous mode");
output.print(" >> describe what you want, I'll build it iteratively"); output.print(" >> describe what you want, I'll build it iteratively");
output.print(""); output.print("");
print!("{}workspace: {}{}\n", print!(
"{}workspace: {}{}\n",
SetForegroundColor(Color::DarkGrey), SetForegroundColor(Color::DarkGrey),
workspace_dir.display(), ResetColor); workspace_dir.display(),
ResetColor
);
output.print(""); output.print("");
output.print("💡 Each input you provide will be added to requirements"); output.print("💡 Each input you provide will be added to requirements");
output.print(" and I'll automatically work on implementing them. You can"); output.print(" and I'll automatically work on implementing them. You can");
@@ -565,7 +597,10 @@ async fn run_accumulative_mode(
if accumulated_requirements.is_empty() { if accumulated_requirements.is_empty() {
output.print("📝 What would you like me to build? (describe your requirements)"); output.print("📝 What would you like me to build? (describe your requirements)");
} else { } else {
output.print(&format!("📝 Turn {} - What's next? (add more requirements or refinements)", turn_number + 1)); output.print(&format!(
"📝 Turn {} - What's next? (add more requirements or refinements)",
turn_number + 1
));
} }
output.print(&format!("{}", "=".repeat(60))); output.print(&format!("{}", "=".repeat(60)));
@@ -628,11 +663,14 @@ async fn run_accumulative_mode(
}; };
// Combine with existing content (README/AGENTS.md) // Combine with existing content (README/AGENTS.md)
let chat_combined_content = match (requirements_context, combined_content.clone()) { let chat_combined_content =
(Some(req_ctx), Some(existing)) => Some(format!("{}\n\n{}", req_ctx, existing)), match (requirements_context, combined_content.clone()) {
(Some(req_ctx), None) => Some(req_ctx), (Some(req_ctx), Some(existing)) => {
(None, existing) => existing, Some(format!("{}\n\n{}", req_ctx, existing))
}; }
(Some(req_ctx), None) => Some(req_ctx),
(None, existing) => existing,
};
// Load configuration // Load configuration
let mut config = Config::load_with_overrides( let mut config = Config::load_with_overrides(
@@ -667,14 +705,24 @@ async fn run_accumulative_mode(
.await?; .await?;
// Run interactive mode // Run interactive mode
run_interactive(agent, cli.show_prompt, cli.show_code, chat_combined_content, &workspace_dir).await?; run_interactive(
agent,
cli.show_prompt,
cli.show_code,
chat_combined_content,
&workspace_dir,
)
.await?;
// After returning from interactive mode, exit // After returning from interactive mode, exit
output.print("\n👋 Goodbye!"); output.print("\n👋 Goodbye!");
break; break;
} }
_ => { _ => {
output.print(&format!("❌ Unknown command: {}. Type /help for available commands.", input)); output.print(&format!(
"❌ Unknown command: {}. Type /help for available commands.",
input
));
continue; continue;
} }
} }
@@ -700,7 +748,10 @@ async fn run_accumulative_mode(
); );
output.print(""); output.print("");
output.print(&format!("📋 Current instructions and requirements (Turn {}):", turn_number)); output.print(&format!(
"📋 Current instructions and requirements (Turn {}):",
turn_number
));
output.print(&format!(" {}", input)); output.print(&format!(" {}", input));
output.print(""); output.print("");
output.print("🚀 Starting autonomous implementation..."); output.print("🚀 Starting autonomous implementation...");
@@ -709,7 +760,7 @@ async fn run_accumulative_mode(
// Create a project with the accumulated requirements // Create a project with the accumulated requirements
let project = Project::new_autonomous_with_requirements( let project = Project::new_autonomous_with_requirements(
workspace_dir.clone(), workspace_dir.clone(),
requirements_doc.clone() requirements_doc.clone(),
)?; )?;
// Ensure workspace exists and enter it // Ensure workspace exists and enter it
@@ -765,8 +816,7 @@ async fn run_accumulative_mode(
} }
}; };
match autonomous_result match autonomous_result {
{
Ok(_) => { Ok(_) => {
output.print(""); output.print("");
output.print("✅ Autonomous run completed"); output.print("✅ Autonomous run completed");
@@ -840,7 +890,9 @@ async fn run_autonomous_machine(
); );
println!("TASK_START"); println!("TASK_START");
let result = agent.execute_task_with_timing(&task, None, false, show_prompt, show_code, true, None).await?; let result = agent
.execute_task_with_timing(&task, None, false, show_prompt, show_code, true, None)
.await?;
println!("AGENT_RESPONSE:"); println!("AGENT_RESPONSE:");
println!("{}", result.response); println!("{}", result.response);
println!("END_AGENT_RESPONSE"); println!("END_AGENT_RESPONSE");
@@ -856,7 +908,6 @@ async fn run_with_console_mode(
project: Project, project: Project,
combined_content: Option<String>, combined_content: Option<String>,
) -> Result<()> { ) -> Result<()> {
// Execute task, autonomous mode, or start interactive mode // Execute task, autonomous mode, or start interactive mode
if cli.autonomous { if cli.autonomous {
// Autonomous mode with coach-player feedback loop // Autonomous mode with coach-player feedback loop
@@ -874,12 +925,27 @@ async fn run_with_console_mode(
// Single-shot mode // Single-shot mode
let output = SimpleOutput::new(); let output = SimpleOutput::new();
let result = agent let result = agent
.execute_task_with_timing(&task, None, false, cli.show_prompt, cli.show_code, true, None) .execute_task_with_timing(
&task,
None,
false,
cli.show_prompt,
cli.show_code,
true,
None,
)
.await?; .await?;
output.print_smart(&result.response); output.print_smart(&result.response);
} else { } else {
// Interactive mode (default) // Interactive mode (default)
run_interactive(agent, cli.show_prompt, cli.show_code, combined_content, project.workspace()).await?; run_interactive(
agent,
cli.show_prompt,
cli.show_code,
combined_content,
project.workspace(),
)
.await?;
} }
Ok(()) Ok(())
@@ -905,7 +971,15 @@ async fn run_with_machine_mode(
} else if let Some(task) = cli.task { } else if let Some(task) = cli.task {
// Single-shot mode // Single-shot mode
let result = agent let result = agent
.execute_task_with_timing(&task, None, false, cli.show_prompt, cli.show_code, true, None) .execute_task_with_timing(
&task,
None,
false,
cli.show_prompt,
cli.show_code,
true,
None,
)
.await?; .await?;
println!("AGENT_RESPONSE:"); println!("AGENT_RESPONSE:");
println!("{}", result.response); println!("{}", result.response);
@@ -943,9 +1017,10 @@ fn read_agents_config(workspace_dir: &Path) -> Option<String> {
let alt_path = workspace_dir.join("agents.md"); let alt_path = workspace_dir.join("agents.md");
if alt_path.exists() { if alt_path.exists() {
match std::fs::read_to_string(&alt_path) { match std::fs::read_to_string(&alt_path) {
Ok(content) => { Ok(content) => Some(format!(
Some(format!("🤖 Agent Configuration (from agents.md):\n\n{}", content)) "🤖 Agent Configuration (from agents.md):\n\n{}",
} content
)),
Err(e) => { Err(e) => {
error!("Failed to read agents.md: {}", e); error!("Failed to read agents.md: {}", e);
None None
@@ -1069,9 +1144,14 @@ async fn run_interactive<W: UiWriter>(
// Display provider and model information // Display provider and model information
match agent.get_provider_info() { match agent.get_provider_info() {
Ok((provider, model)) => { Ok((provider, model)) => {
print!("🔧 {}{}{} | {}{}{}\n", print!(
SetForegroundColor(Color::Cyan), provider, ResetColor, "🔧 {}{}{} | {}{}{}\n",
SetForegroundColor(Color::Yellow), model, ResetColor SetForegroundColor(Color::Cyan),
provider,
ResetColor,
SetForegroundColor(Color::Yellow),
model,
ResetColor
); );
} }
Err(e) => { Err(e) => {
@@ -1086,8 +1166,11 @@ async fn run_interactive<W: UiWriter>(
let has_readme = content.contains("Project README"); let has_readme = content.contains("Project README");
if has_agents { if has_agents {
print!("{}🤖 AGENTS.md configuration loaded{}\n", print!(
SetForegroundColor(Color::DarkGrey), ResetColor); "{}🤖 AGENTS.md configuration loaded{}\n",
SetForegroundColor(Color::DarkGrey),
ResetColor
);
} }
if has_readme { if has_readme {
@@ -1095,17 +1178,22 @@ async fn run_interactive<W: UiWriter>(
let readme_snippet = extract_readme_heading(content) let readme_snippet = extract_readme_heading(content)
.unwrap_or_else(|| "Project documentation loaded".to_string()); .unwrap_or_else(|| "Project documentation loaded".to_string());
print!("{}📚 detected: {}{}\n", print!(
"{}📚 detected: {}{}\n",
SetForegroundColor(Color::DarkGrey), SetForegroundColor(Color::DarkGrey),
readme_snippet, readme_snippet,
ResetColor); ResetColor
);
} }
} }
// Display workspace path // Display workspace path
print!("{}workspace: {}{}\n", print!(
"{}workspace: {}{}\n",
SetForegroundColor(Color::DarkGrey), SetForegroundColor(Color::DarkGrey),
workspace_path.display(), ResetColor); workspace_path.display(),
ResetColor
);
output.print(""); output.print("");
// Initialize rustyline editor with history // Initialize rustyline editor with history
@@ -1190,7 +1278,9 @@ async fn run_interactive<W: UiWriter>(
output.print("📖 Control Commands:"); output.print("📖 Control Commands:");
output.print(" /compact - Trigger auto-summarization (compacts conversation history)"); output.print(" /compact - Trigger auto-summarization (compacts conversation history)");
output.print(" /thinnify - Trigger context thinning (replaces large tool results with file references)"); output.print(" /thinnify - Trigger context thinning (replaces large tool results with file references)");
output.print(" /readme - Reload README.md and AGENTS.md from disk"); output.print(
" /readme - Reload README.md and AGENTS.md from disk",
);
output.print(" /stats - Show detailed context and performance statistics"); output.print(" /stats - Show detailed context and performance statistics");
output.print(" /help - Show this help message"); output.print(" /help - Show this help message");
output.print(" exit/quit - Exit the interactive session"); output.print(" exit/quit - Exit the interactive session");
@@ -1207,7 +1297,10 @@ async fn run_interactive<W: UiWriter>(
output.print("⚠️ Summarization failed"); output.print("⚠️ Summarization failed");
} }
Err(e) => { Err(e) => {
output.print(&format!("❌ Error during summarization: {}", e)); output.print(&format!(
"❌ Error during summarization: {}",
e
));
} }
} }
continue; continue;
@@ -1220,9 +1313,14 @@ async fn run_interactive<W: UiWriter>(
"/readme" => { "/readme" => {
output.print("📚 Reloading README.md and AGENTS.md..."); output.print("📚 Reloading README.md and AGENTS.md...");
match agent.reload_readme() { match agent.reload_readme() {
Ok(true) => output.print("✅ README content reloaded successfully"), Ok(true) => {
Ok(false) => output.print("⚠️ No README was loaded at startup, cannot reload"), output.print("✅ README content reloaded successfully")
Err(e) => output.print(&format!("❌ Error reloading README: {}", e)), }
Ok(false) => output
.print("⚠️ No README was loaded at startup, cannot reload"),
Err(e) => {
output.print(&format!("❌ Error reloading README: {}", e))
}
} }
continue; continue;
} }
@@ -1232,7 +1330,10 @@ async fn run_interactive<W: UiWriter>(
continue; continue;
} }
_ => { _ => {
output.print(&format!("❌ Unknown command: {}. Type /help for available commands.", input)); output.print(&format!(
"❌ Unknown command: {}. Type /help for available commands.",
input
));
continue; continue;
} }
} }
@@ -1421,8 +1522,12 @@ async fn run_interactive_machine(
"/readme" => { "/readme" => {
println!("COMMAND: readme"); println!("COMMAND: readme");
match agent.reload_readme() { match agent.reload_readme() {
Ok(true) => println!("RESULT: README content reloaded successfully"), Ok(true) => {
Ok(false) => println!("RESULT: No README was loaded at startup, cannot reload"), println!("RESULT: README content reloaded successfully")
}
Ok(false) => println!(
"RESULT: No README was loaded at startup, cannot reload"
),
Err(e) => println!("ERROR: {}", e), Err(e) => println!("ERROR: {}", e),
} }
continue; continue;
@@ -1527,7 +1632,10 @@ async fn execute_task_machine(
let delay_ms = 1000 * (2_u64.pow(attempt - 1)); let delay_ms = 1000 * (2_u64.pow(attempt - 1));
let delay = std::time::Duration::from_millis(delay_ms); let delay = std::time::Duration::from_millis(delay_ms);
println!("TIMEOUT: attempt {} of {}, retrying in {:?}", attempt, MAX_TIMEOUT_RETRIES, delay); println!(
"TIMEOUT: attempt {} of {}, retrying in {:?}",
attempt, MAX_TIMEOUT_RETRIES, delay
);
// Wait before retrying // Wait before retrying
tokio::time::sleep(delay).await; tokio::time::sleep(delay).await;
@@ -1594,14 +1702,26 @@ fn display_context_progress<W: UiWriter>(agent: &Agent<W>, _output: &SimpleOutpu
} else if percentage < 60.0 { } else if percentage < 60.0 {
Color::Yellow Color::Yellow
} else if percentage < 80.0 { } else if percentage < 80.0 {
Color::Rgb { r: 255, g: 165, b: 0 } // Orange Color::Rgb {
r: 255,
g: 165,
b: 0,
} // Orange
} else { } else {
Color::Red Color::Red
}; };
// Print with colored dots (using print! directly to handle color codes) // Print with colored dots (using print! directly to handle color codes)
print!("Context: {}{}{}{} {:.0}% ({}/{} tokens)\n", print!(
SetForegroundColor(color), filled_str, empty_str, ResetColor, percentage, context.used_tokens, context.total_tokens); "Context: {}{}{}{} {:.0}% ({}/{} tokens)\n",
SetForegroundColor(color),
filled_str,
empty_str,
ResetColor,
percentage,
context.used_tokens,
context.total_tokens
);
} }
/// Set up the workspace directory for autonomous mode /// Set up the workspace directory for autonomous mode
@@ -1763,35 +1883,53 @@ async fn run_autonomous(
// Load fast-discovery messages before the loop starts (if enabled) // Load fast-discovery messages before the loop starts (if enabled)
let (discovery_messages, discovery_working_dir): (Vec<g3_providers::Message>, Option<String>) = let (discovery_messages, discovery_working_dir): (Vec<g3_providers::Message>, Option<String>) =
if let Some(ref codebase_path) = codebase_fast_start { if let Some(ref codebase_path) = codebase_fast_start {
// Canonicalize the path to ensure it's absolute // Canonicalize the path to ensure it's absolute
let canonical_path = codebase_path.canonicalize().unwrap_or_else(|_| codebase_path.clone()); let canonical_path = codebase_path
let path_str = canonical_path.to_string_lossy(); .canonicalize()
output.print(&format!("🔍 Fast-discovery mode: will explore codebase at {}", path_str)); .unwrap_or_else(|_| codebase_path.clone());
// Get the provider from the agent and use async LLM-based discovery let path_str = canonical_path.to_string_lossy();
match agent.get_provider() { output.print(&format!(
Ok(provider) => { "🔍 Fast-discovery mode: will explore codebase at {}",
// Create a status callback that prints to output path_str
let output_clone = output.clone(); ));
let status_callback: g3_planner::StatusCallback = Box::new(move |msg: &str| { // Get the provider from the agent and use async LLM-based discovery
output_clone.print(msg); match agent.get_provider() {
}); Ok(provider) => {
match g3_planner::get_initial_discovery_messages(&path_str, Some(&requirements), provider, Some(&status_callback)).await { // Create a status callback that prints to output
Ok(messages) => (messages, Some(path_str.to_string())), let output_clone = output.clone();
Err(e) => { let status_callback: g3_planner::StatusCallback = Box::new(move |msg: &str| {
output.print(&format!("⚠️ LLM discovery failed: {}, skipping fast-start", e)); output_clone.print(msg);
(Vec::new(), None) });
match g3_planner::get_initial_discovery_messages(
&path_str,
Some(&requirements),
provider,
Some(&status_callback),
)
.await
{
Ok(messages) => (messages, Some(path_str.to_string())),
Err(e) => {
output.print(&format!(
"⚠️ LLM discovery failed: {}, skipping fast-start",
e
));
(Vec::new(), None)
}
} }
} }
Err(e) => {
output.print(&format!(
"⚠️ Could not get provider: {}, skipping fast-start",
e
));
(Vec::new(), None)
}
} }
Err(e) => { } else {
output.print(&format!("⚠️ Could not get provider: {}, skipping fast-start", e)); (Vec::new(), None)
(Vec::new(), None) };
}
}
} else {
(Vec::new(), None)
};
let has_discovery = !discovery_messages.is_empty(); let has_discovery = !discovery_messages.is_empty();
let mut turn = 1; let mut turn = 1;
@@ -1823,7 +1961,10 @@ async fn run_autonomous(
) )
}; };
output.print(&format!("🎯 Starting player implementation... (elapsed: {})", format_elapsed_time(loop_start.elapsed()))); output.print(&format!(
"🎯 Starting player implementation... (elapsed: {})",
format_elapsed_time(loop_start.elapsed())
));
// Display what feedback the player is receiving // Display what feedback the player is receiving
// If there's no coach feedback on subsequent turns, this is an error // If there's no coach feedback on subsequent turns, this is an error
@@ -1863,7 +2004,9 @@ async fn run_autonomous(
messages: &discovery_messages, messages: &discovery_messages,
fast_start_path: discovery_working_dir.as_deref(), fast_start_path: discovery_working_dir.as_deref(),
}) })
} else { None }, } else {
None
},
) )
.await .await
{ {
@@ -1878,7 +2021,10 @@ async fn run_autonomous(
use g3_core::error_handling::{classify_error, ErrorType, RecoverableError}; use g3_core::error_handling::{classify_error, ErrorType, RecoverableError};
let error_type = classify_error(&e); let error_type = classify_error(&e);
if matches!(error_type, ErrorType::Recoverable(RecoverableError::ContextLengthExceeded)) { if matches!(
error_type,
ErrorType::Recoverable(RecoverableError::ContextLengthExceeded)
) {
output.print(&format!("⚠️ Context length exceeded in player turn: {}", e)); output.print(&format!("⚠️ Context length exceeded in player turn: {}", e));
output.print("📝 Logging error to session and ending current turn..."); output.print("📝 Logging error to session and ending current turn...");
@@ -1924,10 +2070,7 @@ async fn run_autonomous(
output.print("📝 Final Status: 💥 PLAYER PANIC"); output.print("📝 Final Status: 💥 PLAYER PANIC");
output.print("\n📈 Token Usage Statistics:"); output.print("\n📈 Token Usage Statistics:");
output.print(&format!( output.print(&format!(" • Used Tokens: {}", context_window.used_tokens));
" • Used Tokens: {}",
context_window.used_tokens
));
output.print(&format!( output.print(&format!(
" • Total Available: {}", " • Total Available: {}",
context_window.total_tokens context_window.total_tokens
@@ -1954,9 +2097,8 @@ async fn run_autonomous(
)); ));
if _player_retry_count >= MAX_PLAYER_RETRIES { if _player_retry_count >= MAX_PLAYER_RETRIES {
output.print( output
"🔄 Max retries reached for player, marking turn as failed...", .print("🔄 Max retries reached for player, marking turn as failed...");
);
player_failed = true; player_failed = true;
break; // Exit retry loop break; // Exit retry loop
} }
@@ -1973,7 +2115,10 @@ async fn run_autonomous(
)); ));
// Record turn metrics before incrementing // Record turn metrics before incrementing
let turn_duration = turn_start_time.elapsed(); let turn_duration = turn_start_time.elapsed();
let turn_tokens = agent.get_context_window().used_tokens.saturating_sub(turn_start_tokens); let turn_tokens = agent
.get_context_window()
.used_tokens
.saturating_sub(turn_start_tokens);
turn_metrics.push(TurnMetrics { turn_metrics.push(TurnMetrics {
turn_number: turn, turn_number: turn,
tokens_used: turn_tokens, tokens_used: turn_tokens,
@@ -2006,7 +2151,8 @@ async fn run_autonomous(
let ui_writer = ConsoleUiWriter::new(); let ui_writer = ConsoleUiWriter::new();
let mut coach_agent = let mut coach_agent =
Agent::new_autonomous_with_readme_and_quiet(coach_config, ui_writer, None, quiet).await?; Agent::new_autonomous_with_readme_and_quiet(coach_config, ui_writer, None, quiet)
.await?;
// Surface provider info for coach agent // Surface provider info for coach agent
coach_agent.print_provider_banner("Coach"); coach_agent.print_provider_banner("Coach");
@@ -2050,7 +2196,10 @@ Remember: Be clear in your review and concise in your feedback. APPROVE iff the
requirements requirements
); );
output.print(&format!("🎓 Starting coach review... (elapsed: {})", format_elapsed_time(loop_start.elapsed()))); output.print(&format!(
"🎓 Starting coach review... (elapsed: {})",
format_elapsed_time(loop_start.elapsed())
));
// Execute coach task with retry on error // Execute coach task with retry on error
let mut coach_retry_count = 0; let mut coach_retry_count = 0;
@@ -2060,13 +2209,22 @@ Remember: Be clear in your review and concise in your feedback. APPROVE iff the
loop { loop {
match coach_agent match coach_agent
.execute_task_with_timing(&coach_prompt, None, false, show_prompt, show_code, true, .execute_task_with_timing(
&coach_prompt,
None,
false,
show_prompt,
show_code,
true,
if has_discovery { if has_discovery {
Some(DiscoveryOptions { Some(DiscoveryOptions {
messages: &discovery_messages, messages: &discovery_messages,
fast_start_path: discovery_working_dir.as_deref(), fast_start_path: discovery_working_dir.as_deref(),
}) })
} else { None }) } else {
None
},
)
.await .await
{ {
Ok(result) => { Ok(result) => {
@@ -2078,7 +2236,10 @@ Remember: Be clear in your review and concise in your feedback. APPROVE iff the
use g3_core::error_handling::{classify_error, ErrorType, RecoverableError}; use g3_core::error_handling::{classify_error, ErrorType, RecoverableError};
let error_type = classify_error(&e); let error_type = classify_error(&e);
if matches!(error_type, ErrorType::Recoverable(RecoverableError::ContextLengthExceeded)) { if matches!(
error_type,
ErrorType::Recoverable(RecoverableError::ContextLengthExceeded)
) {
output.print(&format!("⚠️ Context length exceeded in coach turn: {}", e)); output.print(&format!("⚠️ Context length exceeded in coach turn: {}", e));
output.print("📝 Logging error to session and ending current turn..."); output.print("📝 Logging error to session and ending current turn...");
@@ -2174,7 +2335,10 @@ Remember: Be clear in your review and concise in your feedback. APPROVE iff the
coach_feedback = "The implementation needs review. Please ensure all requirements are met and the code compiles without errors.".to_string(); coach_feedback = "The implementation needs review. Please ensure all requirements are met and the code compiles without errors.".to_string();
// Record turn metrics before incrementing // Record turn metrics before incrementing
let turn_duration = turn_start_time.elapsed(); let turn_duration = turn_start_time.elapsed();
let turn_tokens = agent.get_context_window().used_tokens.saturating_sub(turn_start_tokens); let turn_tokens = agent
.get_context_window()
.used_tokens
.saturating_sub(turn_start_tokens);
turn_metrics.push(TurnMetrics { turn_metrics.push(TurnMetrics {
turn_number: turn, turn_number: turn,
tokens_used: turn_tokens, tokens_used: turn_tokens,
@@ -2210,7 +2374,10 @@ Remember: Be clear in your review and concise in your feedback. APPROVE iff the
coach_feedback = "The implementation needs review. Please ensure all requirements are met and the code compiles without errors.".to_string(); coach_feedback = "The implementation needs review. Please ensure all requirements are met and the code compiles without errors.".to_string();
// Record turn metrics before incrementing // Record turn metrics before incrementing
let turn_duration = turn_start_time.elapsed(); let turn_duration = turn_start_time.elapsed();
let turn_tokens = agent.get_context_window().used_tokens.saturating_sub(turn_start_tokens); let turn_tokens = agent
.get_context_window()
.used_tokens
.saturating_sub(turn_start_tokens);
turn_metrics.push(TurnMetrics { turn_metrics.push(TurnMetrics {
turn_number: turn, turn_number: turn,
tokens_used: turn_tokens, tokens_used: turn_tokens,
@@ -2241,7 +2408,10 @@ Remember: Be clear in your review and concise in your feedback. APPROVE iff the
coach_feedback = coach_feedback_text; coach_feedback = coach_feedback_text;
// Record turn metrics before incrementing // Record turn metrics before incrementing
let turn_duration = turn_start_time.elapsed(); let turn_duration = turn_start_time.elapsed();
let turn_tokens = agent.get_context_window().used_tokens.saturating_sub(turn_start_tokens); let turn_tokens = agent
.get_context_window()
.used_tokens
.saturating_sub(turn_start_tokens);
turn_metrics.push(TurnMetrics { turn_metrics.push(TurnMetrics {
turn_number: turn, turn_number: turn,
tokens_used: turn_tokens, tokens_used: turn_tokens,
@@ -2296,9 +2466,15 @@ Remember: Be clear in your review and concise in your feedback. APPROVE iff the
output.print(&"=".repeat(60)); output.print(&"=".repeat(60));
if implementation_approved { if implementation_approved {
output.print(&format!("\n🎉 Autonomous mode completed successfully (total loop time: {})", format_elapsed_time(loop_start.elapsed()))); output.print(&format!(
"\n🎉 Autonomous mode completed successfully (total loop time: {})",
format_elapsed_time(loop_start.elapsed())
));
} else { } else {
output.print(&format!("\n🔄 Autonomous mode terminated (max iterations) (total loop time: {})", format_elapsed_time(loop_start.elapsed()))); output.print(&format!(
"\n🔄 Autonomous mode terminated (max iterations) (total loop time: {})",
format_elapsed_time(loop_start.elapsed())
));
} }
Ok(()) Ok(())
+1 -1
View File
@@ -89,7 +89,7 @@ impl UiWriter for MachineUiWriter {
} }
fn wants_full_output(&self) -> bool { fn wants_full_output(&self) -> bool {
true // Machine mode wants complete, untruncated output true // Machine mode wants complete, untruncated output
} }
fn prompt_user_yes_no(&self, message: &str) -> bool { fn prompt_user_yes_no(&self, message: &str) -> bool {
+3 -1
View File
@@ -6,7 +6,9 @@ pub struct SimpleOutput {
impl SimpleOutput { impl SimpleOutput {
pub fn new() -> Self { pub fn new() -> Self {
SimpleOutput { machine_mode: false } SimpleOutput {
machine_mode: false,
}
} }
pub fn new_with_mode(machine_mode: bool) -> Self { pub fn new_with_mode(machine_mode: bool) -> Self {
+22 -11
View File
@@ -48,7 +48,8 @@ impl ConsoleUiWriter {
println!("{}☐ {}", indent, task); println!("{}☐ {}", indent, task);
} else if trimmed.starts_with("- [x]") || trimmed.starts_with("- [X]") { } else if trimmed.starts_with("- [x]") || trimmed.starts_with("- [X]") {
// Completed task // Completed task
let task = trimmed.strip_prefix("- [x]") let task = trimmed
.strip_prefix("- [x]")
.or_else(|| trimmed.strip_prefix("- [X]")) .or_else(|| trimmed.strip_prefix("- [X]"))
.unwrap_or(trimmed) .unwrap_or(trimmed)
.trim(); .trim();
@@ -108,9 +109,9 @@ impl UiWriter for ConsoleUiWriter {
// Flash animation: print with bright background, then normal // Flash animation: print with bright background, then normal
let frames = vec![ let frames = vec![
"\x1b[1;97;46m", // Frame 1: Bold white on cyan background "\x1b[1;97;46m", // Frame 1: Bold white on cyan background
"\x1b[1;97;42m", // Frame 2: Bold white on green background "\x1b[1;97;42m", // Frame 2: Bold white on green background
"\x1b[1;96;40m", // Frame 3: Bold cyan on black background "\x1b[1;96;40m", // Frame 3: Bold cyan on black background
]; ];
println!(); println!();
@@ -143,8 +144,7 @@ impl UiWriter for ConsoleUiWriter {
*self.in_todo_tool.lock().unwrap() = is_todo; *self.in_todo_tool.lock().unwrap() = is_todo;
// For todo tools, we'll skip the normal header and print a custom one later // For todo tools, we'll skip the normal header and print a custom one later
if is_todo { if is_todo {}
}
} }
fn print_tool_arg(&self, key: &str, value: &str) { fn print_tool_arg(&self, key: &str, value: &str) {
@@ -192,7 +192,8 @@ impl UiWriter for ConsoleUiWriter {
// Truncate long values for display // Truncate long values for display
let display_value = if first_line.len() > 80 { let display_value = if first_line.len() > 80 {
// Use char_indices to safely truncate at character boundary // Use char_indices to safely truncate at character boundary
let truncate_at = first_line.char_indices() let truncate_at = first_line
.char_indices()
.nth(77) .nth(77)
.map(|(i, _)| i) .map(|(i, _)| i)
.unwrap_or(first_line.len()); .unwrap_or(first_line.len());
@@ -208,8 +209,16 @@ impl UiWriter for ConsoleUiWriter {
let has_end = args.iter().any(|(k, _)| k == "end"); let has_end = args.iter().any(|(k, _)| k == "end");
if has_start || has_end { if has_start || has_end {
let start_val = args.iter().find(|(k, _)| k == "start").map(|(_, v)| v.as_str()).unwrap_or("0"); let start_val = args
let end_val = args.iter().find(|(k, _)| k == "end").map(|(_, v)| v.as_str()).unwrap_or("end"); .iter()
.find(|(k, _)| k == "start")
.map(|(_, v)| v.as_str())
.unwrap_or("0");
let end_val = args
.iter()
.find(|(k, _)| k == "end")
.map(|(_, v)| v.as_str())
.unwrap_or("end");
format!(" [{}..{}]", start_val, end_val) format!(" [{}..{}]", start_val, end_val)
} else { } else {
String::new() String::new()
@@ -219,7 +228,10 @@ impl UiWriter for ConsoleUiWriter {
}; };
// Print with bold green tool name, purple (non-bold) for pipe and args // Print with bold green tool name, purple (non-bold) for pipe and args
println!("┌─\x1b[1;32m {}\x1b[0m\x1b[35m | {}{}\x1b[0m", tool_name, display_value, header_suffix); println!(
"┌─\x1b[1;32m {}\x1b[0m\x1b[35m | {}{}\x1b[0m",
tool_name, display_value, header_suffix
);
} else { } else {
// Print with bold green formatting using ANSI escape codes // Print with bold green formatting using ANSI escape codes
println!("┌─\x1b[1;32m {}\x1b[0m", tool_name); println!("┌─\x1b[1;32m {}\x1b[0m", tool_name);
@@ -379,4 +391,3 @@ impl UiWriter for ConsoleUiWriter {
} }
} }
} }
+25 -9
View File
@@ -34,12 +34,17 @@ fn main() {
.expect("Failed to find .build/release directory"); .expect("Failed to find .build/release directory");
// Copy the dylib to the output directory so it can be found at runtime // Copy the dylib to the output directory so it can be found at runtime
let target_dir = manifest_dir.parent().unwrap().parent().unwrap().join("target"); let target_dir = manifest_dir
.parent()
.unwrap()
.parent()
.unwrap()
.join("target");
let profile = env::var("PROFILE").unwrap_or_else(|_| "debug".to_string()); let profile = env::var("PROFILE").unwrap_or_else(|_| "debug".to_string());
// Determine the actual target directory (could be llvm-cov-target or regular target) // Determine the actual target directory (could be llvm-cov-target or regular target)
let target_dir_name = env::var("CARGO_TARGET_DIR") let target_dir_name =
.unwrap_or_else(|_| target_dir.to_string_lossy().to_string()); env::var("CARGO_TARGET_DIR").unwrap_or_else(|_| target_dir.to_string_lossy().to_string());
let actual_target_dir = PathBuf::from(&target_dir_name); let actual_target_dir = PathBuf::from(&target_dir_name);
let output_dir = actual_target_dir.join(&profile); let output_dir = actual_target_dir.join(&profile);
@@ -47,13 +52,21 @@ fn main() {
let dylib_dst = output_dir.join("libVisionBridge.dylib"); let dylib_dst = output_dir.join("libVisionBridge.dylib");
// Create output directory if it doesn't exist // Create output directory if it doesn't exist
std::fs::create_dir_all(&output_dir) std::fs::create_dir_all(&output_dir).expect(&format!(
.expect(&format!("Failed to create output directory {}", output_dir.display())); "Failed to create output directory {}",
output_dir.display()
));
std::fs::copy(&dylib_src, &dylib_dst) std::fs::copy(&dylib_src, &dylib_dst).expect(&format!(
.expect(&format!("Failed to copy dylib from {} to {}", dylib_src.display(), dylib_dst.display())); "Failed to copy dylib from {} to {}",
dylib_src.display(),
dylib_dst.display()
));
println!("cargo:warning=Copied libVisionBridge.dylib to {}", dylib_dst.display()); println!(
"cargo:warning=Copied libVisionBridge.dylib to {}",
dylib_dst.display()
);
// Add rpath so the dylib can be found at runtime // Add rpath so the dylib can be found at runtime
println!("cargo:rustc-link-arg=-Wl,-rpath,@executable_path"); println!("cargo:rustc-link-arg=-Wl,-rpath,@executable_path");
@@ -68,5 +81,8 @@ fn main() {
println!("cargo:rustc-link-lib=framework=CoreGraphics"); println!("cargo:rustc-link-lib=framework=CoreGraphics");
println!("cargo:rustc-link-lib=framework=CoreImage"); println!("cargo:rustc-link-lib=framework=CoreImage");
println!("cargo:warning=VisionBridge built successfully at {}", lib_path.display()); println!(
"cargo:warning=VisionBridge built successfully at {}",
lib_path.display()
);
} }
@@ -23,14 +23,23 @@ fn main() {
println!("\nRow alignment:"); println!("\nRow alignment:");
println!(" Actual bytes per row: {}", bytes_per_row); println!(" Actual bytes per row: {}", bytes_per_row);
println!(" Expected (width * 4): {}", expected_bytes_per_row); println!(" Expected (width * 4): {}", expected_bytes_per_row);
println!(" Padding per row: {}", bytes_per_row - expected_bytes_per_row); println!(
" Padding per row: {}",
bytes_per_row - expected_bytes_per_row
);
// Sample some pixels from different locations // Sample some pixels from different locations
println!("\nFirst 3 pixels (raw bytes):"); println!("\nFirst 3 pixels (raw bytes):");
for i in 0..3 { for i in 0..3 {
let offset = i * 4; let offset = i * 4;
println!(" Pixel {}: [{:3}, {:3}, {:3}, {:3}]", println!(
i, data[offset], data[offset+1], data[offset+2], data[offset+3]); " Pixel {}: [{:3}, {:3}, {:3}, {:3}]",
i,
data[offset],
data[offset + 1],
data[offset + 2],
data[offset + 3]
);
} }
// Check a pixel from the middle // Check a pixel from the middle
@@ -40,7 +49,12 @@ fn main() {
println!("\nMiddle pixel (row {}, col {}):", mid_row, mid_col); println!("\nMiddle pixel (row {}, col {}):", mid_row, mid_col);
println!(" Offset: {}", mid_offset); println!(" Offset: {}", mid_offset);
if mid_offset + 3 < data.len() as usize { if mid_offset + 3 < data.len() as usize {
println!(" Bytes: [{:3}, {:3}, {:3}, {:3}]", println!(
data[mid_offset], data[mid_offset+1], data[mid_offset+2], data[mid_offset+3]); " Bytes: [{:3}, {:3}, {:3}, {:3}]",
data[mid_offset],
data[mid_offset + 1],
data[mid_offset + 2],
data[mid_offset + 3]
);
} }
} }
@@ -1,7 +1,9 @@
use core_graphics::window::{kCGWindowListOptionOnScreenOnly, kCGNullWindowID, CGWindowListCopyWindowInfo}; use core_foundation::base::{TCFType, ToVoid};
use core_foundation::dictionary::CFDictionary; use core_foundation::dictionary::CFDictionary;
use core_foundation::string::CFString; use core_foundation::string::CFString;
use core_foundation::base::{TCFType, ToVoid}; use core_graphics::window::{
kCGNullWindowID, kCGWindowListOptionOnScreenOnly, CGWindowListCopyWindowInfo,
};
fn main() { fn main() {
println!("Listing all on-screen windows..."); println!("Listing all on-screen windows...");
@@ -9,13 +11,14 @@ fn main() {
println!("{}", "-".repeat(80)); println!("{}", "-".repeat(80));
unsafe { unsafe {
let window_list = CGWindowListCopyWindowInfo( let window_list =
kCGWindowListOptionOnScreenOnly, CGWindowListCopyWindowInfo(kCGWindowListOptionOnScreenOnly, kCGNullWindowID);
kCGNullWindowID
);
let count = core_foundation::array::CFArray::<CFDictionary>::wrap_under_create_rule(window_list).len(); let count =
let array = core_foundation::array::CFArray::<CFDictionary>::wrap_under_create_rule(window_list); core_foundation::array::CFArray::<CFDictionary>::wrap_under_create_rule(window_list)
.len();
let array =
core_foundation::array::CFArray::<CFDictionary>::wrap_under_create_rule(window_list);
for i in 0..count { for i in 0..count {
let dict = array.get(i).unwrap(); let dict = array.get(i).unwrap();
@@ -23,7 +26,8 @@ fn main() {
// Get window ID // Get window ID
let window_id_key = CFString::from_static_string("kCGWindowNumber"); let window_id_key = CFString::from_static_string("kCGWindowNumber");
let window_id: i64 = if let Some(value) = dict.find(window_id_key.to_void()) { let window_id: i64 = if let Some(value) = dict.find(window_id_key.to_void()) {
let num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*value as *const _); let num: core_foundation::number::CFNumber =
TCFType::wrap_under_get_rule(*value as *const _);
num.to_i64().unwrap_or(0) num.to_i64().unwrap_or(0)
} else { } else {
0 0
@@ -1,6 +1,6 @@
use g3_computer_control::SafariDriver;
use g3_computer_control::webdriver::WebDriverController;
use anyhow::Result; use anyhow::Result;
use g3_computer_control::webdriver::WebDriverController;
use g3_computer_control::SafariDriver;
#[tokio::main] #[tokio::main]
async fn main() -> Result<()> { async fn main() -> Result<()> {
@@ -47,7 +47,9 @@ async fn main() -> Result<()> {
// Execute JavaScript // Execute JavaScript
println!("Executing JavaScript..."); println!("Executing JavaScript...");
let result = driver.execute_script("return document.title", vec![]).await?; let result = driver
.execute_script("return document.title", vec![])
.await?;
println!("JS result: {:?}\n", result); println!("JS result: {:?}\n", result);
// Take a screenshot // Take a screenshot
@@ -6,7 +6,10 @@ async fn main() {
let controller = create_controller().expect("Failed to create controller"); let controller = create_controller().expect("Failed to create controller");
match controller.take_screenshot("/tmp/test_with_prompt.png", None, None).await { match controller
.take_screenshot("/tmp/test_with_prompt.png", None, None)
.await
{
Ok(_) => { Ok(_) => {
println!("\n✅ Screenshot saved to /tmp/test_with_prompt.png"); println!("\n✅ Screenshot saved to /tmp/test_with_prompt.png");
println!("Opening screenshot..."); println!("Opening screenshot...");
@@ -22,7 +22,11 @@ fn main() {
// Check file exists and size // Check file exists and size
if let Ok(metadata) = std::fs::metadata(path) { if let Ok(metadata) = std::fs::metadata(path) {
println!("File size: {} bytes ({:.1} MB)", metadata.len(), metadata.len() as f64 / 1_000_000.0); println!(
"File size: {} bytes ({:.1} MB)",
metadata.len(),
metadata.len() as f64 / 1_000_000.0
);
} }
// Open it // Open it
@@ -11,9 +11,15 @@ fn main() {
let data = image.data(); let data = image.data();
println!("Testing screenshot fix..."); println!("Testing screenshot fix...");
println!("Image: {}x{}, bytes_per_row: {}", width, height, bytes_per_row); println!(
"Image: {}x{}, bytes_per_row: {}",
width, height, bytes_per_row
);
println!("Expected bytes per row: {}", width * 4); println!("Expected bytes per row: {}", width * 4);
println!("Padding per row: {} bytes", bytes_per_row - (width as usize * 4)); println!(
"Padding per row: {} bytes",
bytes_per_row - (width as usize * 4)
);
// OLD METHOD (broken) - treating data as continuous // OLD METHOD (broken) - treating data as continuous
println!("\n=== OLD METHOD (BROKEN) ==="); println!("\n=== OLD METHOD (BROKEN) ===");
@@ -48,7 +54,11 @@ fn main() {
let crop_size = 200; let crop_size = 200;
// Old method crop // Old method crop
let old_crop: Vec<u8> = old_rgba.iter().take((crop_size * crop_size * 4) as usize).copied().collect(); let old_crop: Vec<u8> = old_rgba
.iter()
.take((crop_size * crop_size * 4) as usize)
.copied()
.collect();
if let Some(old_img) = ImageBuffer::from_raw(crop_size, crop_size, old_crop) { if let Some(old_img) = ImageBuffer::from_raw(crop_size, crop_size, old_crop) {
let old_img: RgbaImage = old_img; let old_img: RgbaImage = old_img;
old_img.save("/tmp/screenshot_old_method.png").unwrap(); old_img.save("/tmp/screenshot_old_method.png").unwrap();
@@ -56,7 +66,11 @@ fn main() {
} }
// New method crop // New method crop
let new_crop: Vec<u8> = new_rgba.iter().take((crop_size * crop_size * 4) as usize).copied().collect(); let new_crop: Vec<u8> = new_rgba
.iter()
.take((crop_size * crop_size * 4) as usize)
.copied()
.collect();
if let Some(new_img) = ImageBuffer::from_raw(crop_size, crop_size, new_crop) { if let Some(new_img) = ImageBuffer::from_raw(crop_size, crop_size, new_crop) {
let new_img: RgbaImage = new_img; let new_img: RgbaImage = new_img;
new_img.save("/tmp/screenshot_new_method.png").unwrap(); new_img.save("/tmp/screenshot_new_method.png").unwrap();
@@ -1,5 +1,5 @@
use g3_computer_control::ocr::{OCREngine, DefaultOCR};
use anyhow::Result; use anyhow::Result;
use g3_computer_control::ocr::{DefaultOCR, OCREngine};
#[tokio::main] #[tokio::main]
async fn main() -> Result<()> { async fn main() -> Result<()> {
@@ -47,7 +47,10 @@ async fn main() -> Result<()> {
println!("⚠️ No text found in image"); println!("⚠️ No text found in image");
} else { } else {
println!(" Top 20 results:"); println!(" Top 20 results:");
println!(" {:<4} {:<40} {:<15} {:<12} {:<8}", "#", "Text", "Position", "Size", "Conf"); println!(
" {:<4} {:<40} {:<15} {:<12} {:<8}",
"#", "Text", "Position", "Size", "Conf"
);
println!(" {}", "-".repeat(85)); println!(" {}", "-".repeat(85));
for (i, loc) in locations.iter().take(20).enumerate() { for (i, loc) in locations.iter().take(20).enumerate() {
@@ -57,7 +60,8 @@ async fn main() -> Result<()> {
loc.text.clone() loc.text.clone()
}; };
println!(" {:<4} {:<40} ({:>4},{:>4}) {:>4}x{:<4} {:.2}", println!(
" {:<4} {:<40} ({:>4},{:>4}) {:>4}x{:<4} {:.2}",
i + 1, i + 1,
text, text,
loc.x, loc.x,
@@ -76,7 +80,10 @@ async fn main() -> Result<()> {
println!("\n📈 Performance:"); println!("\n📈 Performance:");
println!(" OCR Speed: {:.3}s", duration.as_secs_f64()); println!(" OCR Speed: {:.3}s", duration.as_secs_f64());
println!(" Text elements: {}", locations.len()); println!(" Text elements: {}", locations.len());
println!(" Avg per element: {:.1}ms", duration.as_millis() as f64 / locations.len() as f64); println!(
" Avg per element: {:.1}ms",
duration.as_millis() as f64 / locations.len() as f64
);
} }
println!("\n✅ Test complete!"); println!("\n✅ Test complete!");
@@ -8,10 +8,15 @@ async fn main() {
// Test 1: Capture iTerm2 window // Test 1: Capture iTerm2 window
println!("\n1. Capturing iTerm2 window..."); println!("\n1. Capturing iTerm2 window...");
match controller.take_screenshot("/tmp/iterm_window.png", None, Some("iTerm2")).await { match controller
.take_screenshot("/tmp/iterm_window.png", None, Some("iTerm2"))
.await
{
Ok(_) => { Ok(_) => {
println!(" ✅ iTerm2 window captured to /tmp/iterm_window.png"); println!(" ✅ iTerm2 window captured to /tmp/iterm_window.png");
let _ = std::process::Command::new("open").arg("/tmp/iterm_window.png").spawn(); let _ = std::process::Command::new("open")
.arg("/tmp/iterm_window.png")
.spawn();
} }
Err(e) => println!(" ❌ Failed: {}", e), Err(e) => println!(" ❌ Failed: {}", e),
} }
@@ -21,10 +26,15 @@ async fn main() {
// Test 2: Full screen capture for comparison // Test 2: Full screen capture for comparison
println!("\n2. Capturing full screen for comparison..."); println!("\n2. Capturing full screen for comparison...");
match controller.take_screenshot("/tmp/fullscreen.png", None, None).await { match controller
.take_screenshot("/tmp/fullscreen.png", None, None)
.await
{
Ok(_) => { Ok(_) => {
println!(" ✅ Full screen captured to /tmp/fullscreen.png"); println!(" ✅ Full screen captured to /tmp/fullscreen.png");
let _ = std::process::Command::new("open").arg("/tmp/fullscreen.png").spawn(); let _ = std::process::Command::new("open")
.arg("/tmp/fullscreen.png")
.spawn();
} }
Err(e) => println!(" ❌ Failed: {}", e), Err(e) => println!(" ❌ Failed: {}", e),
} }
+17 -8
View File
@@ -1,17 +1,17 @@
// Suppress warnings from objc crate macros // Suppress warnings from objc crate macros
#![allow(unexpected_cfgs)] #![allow(unexpected_cfgs)]
pub mod types;
pub mod platform;
pub mod ocr;
pub mod webdriver;
pub mod macax; pub mod macax;
pub mod ocr;
pub mod platform;
pub mod types;
pub mod webdriver;
// Re-export webdriver types for convenience // Re-export webdriver types for convenience
pub use webdriver::{WebDriverController, WebElement, safari::SafariDriver}; pub use webdriver::{safari::SafariDriver, WebDriverController, WebElement};
// Re-export macax types for convenience // Re-export macax types for convenience
pub use macax::{MacAxController, AXElement, AXApplication}; pub use macax::{AXApplication, AXElement, MacAxController};
use anyhow::Result; use anyhow::Result;
use async_trait::async_trait; use async_trait::async_trait;
@@ -20,13 +20,22 @@ use types::*;
#[async_trait] #[async_trait]
pub trait ComputerController: Send + Sync { pub trait ComputerController: Send + Sync {
// Screen capture // Screen capture
async fn take_screenshot(&self, path: &str, region: Option<Rect>, window_id: Option<&str>) -> Result<()>; async fn take_screenshot(
&self,
path: &str,
region: Option<Rect>,
window_id: Option<&str>,
) -> Result<()>;
// OCR operations // OCR operations
async fn extract_text_from_screen(&self, region: Rect, window_id: &str) -> Result<String>; async fn extract_text_from_screen(&self, region: Rect, window_id: &str) -> Result<String>;
async fn extract_text_from_image(&self, path: &str) -> Result<String>; async fn extract_text_from_image(&self, path: &str) -> Result<String>;
async fn extract_text_with_locations(&self, path: &str) -> Result<Vec<TextLocation>>; async fn extract_text_with_locations(&self, path: &str) -> Result<Vec<TextLocation>>;
async fn find_text_in_app(&self, app_name: &str, search_text: &str) -> Result<Option<TextLocation>>; async fn find_text_in_app(
&self,
app_name: &str,
search_text: &str,
) -> Result<Option<TextLocation>>;
// Mouse operations // Mouse operations
fn move_mouse(&self, x: i32, y: i32) -> Result<()>; fn move_mouse(&self, x: i32, y: i32) -> Result<()>;
@@ -3,7 +3,9 @@ use anyhow::{Context, Result};
use std::collections::HashMap; use std::collections::HashMap;
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
use accessibility::{AXUIElement, AXUIElementAttributes, ElementFinder, TreeVisitor, TreeWalker, TreeWalkerFlow}; use accessibility::{
AXUIElement, AXUIElementAttributes, ElementFinder, TreeVisitor, TreeWalker, TreeWalkerFlow,
};
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
use core_foundation::base::TCFType; use core_foundation::base::TCFType;
@@ -99,7 +101,9 @@ impl MacAxController {
// Skip background-only apps // Skip background-only apps
let activation_policy: i64 = msg_send![app, activationPolicy]; let activation_policy: i64 = msg_send![app, activationPolicy];
if activation_policy == NSApplicationActivationPolicy::NSApplicationActivationPolicyRegular as i64 { if activation_policy
== NSApplicationActivationPolicy::NSApplicationActivationPolicyRegular as i64
{
apps.push(AXApplication { apps.push(AXApplication {
name, name,
bundle_id, bundle_id,
@@ -263,16 +267,17 @@ impl MacAxController {
let indent = " ".repeat(depth); let indent = " ".repeat(depth);
// Get role // Get role
let role = element.role().ok().map(|s| s.to_string()) let role = element
.role()
.ok()
.map(|s| s.to_string())
.unwrap_or_else(|| "Unknown".to_string()); .unwrap_or_else(|| "Unknown".to_string());
// Get title // Get title
let title = element.title().ok() let title = element.title().ok().map(|s| s.to_string());
.map(|s| s.to_string());
// Get identifier // Get identifier
let identifier = element.identifier().ok() let identifier = element.identifier().ok().map(|s| s.to_string());
.map(|s| s.to_string());
// Format output // Format output
output.push_str(&format!("{}Role: {}", indent, role)); output.push_str(&format!("{}Role: {}", indent, role));
@@ -352,9 +357,7 @@ impl MacAxController {
&app_element, &app_element,
move |element| { move |element| {
// Check role // Check role
let elem_role = element.role() let elem_role = element.role().ok().map(|s| s.to_string());
.ok()
.map(|s| s.to_string());
if let Some(r) = elem_role { if let Some(r) = elem_role {
if !r.contains(&role_str) { if !r.contains(&role_str) {
@@ -366,9 +369,7 @@ impl MacAxController {
// Check title if specified // Check title if specified
if let Some(ref title_filter) = title_str { if let Some(ref title_filter) = title_str {
let elem_title = element.title() let elem_title = element.title().ok().map(|s| s.to_string());
.ok()
.map(|s| s.to_string());
if let Some(t) = elem_title { if let Some(t) = elem_title {
if !t.contains(title_filter) { if !t.contains(title_filter) {
@@ -381,9 +382,7 @@ impl MacAxController {
// Check identifier if specified // Check identifier if specified
if let Some(ref id_filter) = identifier_str { if let Some(ref id_filter) = identifier_str {
let elem_id = element.identifier() let elem_id = element.identifier().ok().map(|s| s.to_string());
.ok()
.map(|s| s.to_string());
if let Some(id) = elem_id { if let Some(id) = elem_id {
if !id.contains(id_filter) { if !id.contains(id_filter) {
@@ -448,7 +447,8 @@ impl MacAxController {
// Set the value - convert CFString to CFType // Set the value - convert CFString to CFType
let cf_value = CFString::new(value); let cf_value = CFString::new(value);
element.set_value(cf_value.as_CFType()) element
.set_value(cf_value.as_CFType())
.map_err(|e| anyhow::anyhow!("Failed to set value: {:?}", e))?; .map_err(|e| anyhow::anyhow!("Failed to set value: {:?}", e))?;
Ok(()) Ok(())
@@ -478,7 +478,8 @@ impl MacAxController {
let element = self.find_element(app_name, role, title, identifier)?; let element = self.find_element(app_name, role, title, identifier)?;
// Get the value // Get the value
let value_type = element.value() let value_type = element
.value()
.map_err(|e| anyhow::anyhow!("Failed to get value: {:?}", e))?; .map_err(|e| anyhow::anyhow!("Failed to get value: {:?}", e))?;
// Try to downcast to CFString // Try to downcast to CFString
@@ -579,7 +580,8 @@ impl MacAxController {
use core_foundation::boolean::CFBoolean; use core_foundation::boolean::CFBoolean;
let cf_true = CFBoolean::true_value(); let cf_true = CFBoolean::true_value();
element.set_attribute(&accessibility::AXAttribute::focused(), cf_true) element
.set_attribute(&accessibility::AXAttribute::focused(), cf_true)
.map_err(|e| anyhow::anyhow!("Failed to focus element: {:?}", e))?; .map_err(|e| anyhow::anyhow!("Failed to focus element: {:?}", e))?;
Ok(()) Ok(())
@@ -587,15 +589,8 @@ impl MacAxController {
/// Press a keyboard shortcut /// Press a keyboard shortcut
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
pub fn press_key( pub fn press_key(&self, app_name: &str, key: &str, modifiers: Vec<&str>) -> Result<()> {
&self, use core_graphics::event::{CGEvent, CGEventFlags, CGEventTapLocation};
app_name: &str,
key: &str,
modifiers: Vec<&str>,
) -> Result<()> {
use core_graphics::event::{
CGEvent, CGEventFlags, CGEventTapLocation,
};
use core_graphics::event_source::{CGEventSource, CGEventSourceStateID}; use core_graphics::event_source::{CGEventSource, CGEventSourceStateID};
// First, make sure the app is active // First, make sure the app is active
@@ -605,8 +600,8 @@ impl MacAxController {
std::thread::sleep(std::time::Duration::from_millis(100)); std::thread::sleep(std::time::Duration::from_millis(100));
// Map key string to key code // Map key string to key code
let key_code = Self::key_to_keycode(key) let key_code =
.ok_or_else(|| anyhow::anyhow!("Unknown key: {}", key))?; Self::key_to_keycode(key).ok_or_else(|| anyhow::anyhow!("Unknown key: {}", key))?;
// Map modifiers to flags // Map modifiers to flags
let mut flags = CGEventFlags::CGEventFlagNull; let mut flags = CGEventFlags::CGEventFlagNull;
@@ -622,16 +617,19 @@ impl MacAxController {
// Create event source // Create event source
let source = CGEventSource::new(CGEventSourceStateID::HIDSystemState) let source = CGEventSource::new(CGEventSourceStateID::HIDSystemState)
.ok().context("Failed to create event source")?; .ok()
.context("Failed to create event source")?;
// Create key down event // Create key down event
let key_down = CGEvent::new_keyboard_event(source.clone(), key_code, true) let key_down = CGEvent::new_keyboard_event(source.clone(), key_code, true)
.ok().context("Failed to create key down event")?; .ok()
.context("Failed to create key down event")?;
key_down.set_flags(flags); key_down.set_flags(flags);
// Create key up event // Create key up event
let key_up = CGEvent::new_keyboard_event(source, key_code, false) let key_up = CGEvent::new_keyboard_event(source, key_code, false)
.ok().context("Failed to create key up event")?; .ok()
.context("Failed to create key up event")?;
key_up.set_flags(flags); key_up.set_flags(flags);
// Post events // Post events
@@ -643,12 +641,7 @@ impl MacAxController {
} }
#[cfg(not(target_os = "macos"))] #[cfg(not(target_os = "macos"))]
pub fn press_key( pub fn press_key(&self, _app_name: &str, _key: &str, _modifiers: Vec<&str>) -> Result<()> {
&self,
_app_name: &str,
_key: &str,
_modifiers: Vec<&str>,
) -> Result<()> {
anyhow::bail!("Not supported on this platform") anyhow::bail!("Not supported on this platform")
} }
@@ -749,52 +742,45 @@ impl<'a> TreeVisitor for ElementCollector<'a> {
} }
// Get element properties // Get element properties
let role = element.role() let role = element
.role()
.ok() .ok()
.map(|s| s.to_string()) .map(|s| s.to_string())
.unwrap_or_else(|| "Unknown".to_string()); .unwrap_or_else(|| "Unknown".to_string());
let title = element.title() let title = element.title().ok().map(|s| s.to_string());
.ok()
.map(|s| s.to_string());
let identifier = element.identifier() let identifier = element.identifier().ok().map(|s| s.to_string());
.ok()
.map(|s| s.to_string());
// Check if this element matches the filters // Check if this element matches the filters
let role_matches = self.role_filter.as_ref().map_or(true, |r| role.contains(r)); let role_matches = self.role_filter.as_ref().map_or(true, |r| role.contains(r));
let title_matches = self.title_filter.as_ref().map_or(true, |t| { let title_matches = self.title_filter.as_ref().map_or(true, |t| {
title.as_ref().map_or(false, |title_str| title_str.contains(t)) title
.as_ref()
.map_or(false, |title_str| title_str.contains(t))
}); });
let identifier_matches = self.identifier_filter.as_ref().map_or(true, |id| { let identifier_matches = self.identifier_filter.as_ref().map_or(true, |id| {
identifier.as_ref().map_or(false, |id_str| id_str.contains(id)) identifier
.as_ref()
.map_or(false, |id_str| id_str.contains(id))
}); });
if role_matches && title_matches && identifier_matches { if role_matches && title_matches && identifier_matches {
// Get additional properties // Get additional properties
let value = element.value() let value = element
.value()
.ok() .ok()
.and_then(|v| { .and_then(|v| v.downcast::<CFString>().map(|s| s.to_string()));
v.downcast::<CFString>().map(|s| s.to_string())
});
let label = element.description() let label = element.description().ok().map(|s| s.to_string());
.ok()
.map(|s| s.to_string());
let enabled = element.enabled() let enabled = element.enabled().ok().map(|b| b.into()).unwrap_or(false);
.ok()
.map(|b| b.into())
.unwrap_or(false);
let focused = element.focused() let focused = element.focused().ok().map(|b| b.into()).unwrap_or(false);
.ok()
.map(|b| b.into())
.unwrap_or(false);
// Count children // Count children
let children_count = element.children() let children_count = element
.children()
.ok() .ok()
.map(|arr| arr.len() as usize) .map(|arr| arr.len() as usize)
.unwrap_or(0); .unwrap_or(0);
@@ -14,12 +14,14 @@ impl TesseractOCR {
.output(); .output();
if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() { if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() {
anyhow::bail!("Tesseract OCR is not installed on your system.\n\n\ anyhow::bail!(
"Tesseract OCR is not installed on your system.\n\n\
To install tesseract:\n macOS: brew install tesseract\n \ To install tesseract:\n macOS: brew install tesseract\n \
Linux: sudo apt-get install tesseract-ocr (Ubuntu/Debian)\n \ Linux: sudo apt-get install tesseract-ocr (Ubuntu/Debian)\n \
sudo yum install tesseract (RHEL/CentOS)\n \ sudo yum install tesseract (RHEL/CentOS)\n \
Windows: Download from https://github.com/UB-Mannheim/tesseract/wiki\n\n\ Windows: Download from https://github.com/UB-Mannheim/tesseract/wiki\n\n\
After installation, restart your terminal and try again."); After installation, restart your terminal and try again."
);
} }
Ok(Self) Ok(Self)
@@ -38,7 +40,10 @@ impl OCREngine for TesseractOCR {
.map_err(|e| anyhow::anyhow!("Failed to run tesseract: {}", e))?; .map_err(|e| anyhow::anyhow!("Failed to run tesseract: {}", e))?;
if !output.status.success() { if !output.status.success() {
anyhow::bail!("Tesseract failed: {}", String::from_utf8_lossy(&output.stderr)); anyhow::bail!(
"Tesseract failed: {}",
String::from_utf8_lossy(&output.stderr)
);
} }
let tsv_text = String::from_utf8_lossy(&output.stdout); let tsv_text = String::from_utf8_lossy(&output.stdout);
@@ -46,7 +51,9 @@ impl OCREngine for TesseractOCR {
// Parse TSV output (skip header line) // Parse TSV output (skip header line)
for (i, line) in tsv_text.lines().enumerate() { for (i, line) in tsv_text.lines().enumerate() {
if i == 0 { continue; } // Skip header if i == 0 {
continue;
} // Skip header
let parts: Vec<&str> = line.split('\t').collect(); let parts: Vec<&str> = line.split('\t').collect();
if parts.len() >= 12 { if parts.len() >= 12 {
+3 -6
View File
@@ -1,6 +1,6 @@
use super::OCREngine; use super::OCREngine;
use crate::types::TextLocation; use crate::types::TextLocation;
use anyhow::{Result, Context}; use anyhow::{Context, Result};
use async_trait::async_trait; use async_trait::async_trait;
use std::ffi::{CStr, CString}; use std::ffi::{CStr, CString};
use std::os::raw::{c_char, c_float, c_uint}; use std::os::raw::{c_char, c_float, c_uint};
@@ -41,8 +41,7 @@ impl AppleVisionOCR {
impl OCREngine for AppleVisionOCR { impl OCREngine for AppleVisionOCR {
async fn extract_text_with_locations(&self, path: &str) -> Result<Vec<TextLocation>> { async fn extract_text_with_locations(&self, path: &str) -> Result<Vec<TextLocation>> {
// Convert path to C string // Convert path to C string
let c_path = CString::new(path) let c_path = CString::new(path).context("Failed to convert path to C string")?;
.context("Failed to convert path to C string")?;
let mut boxes_ptr: *mut std::ffi::c_void = std::ptr::null_mut(); let mut boxes_ptr: *mut std::ffi::c_void = std::ptr::null_mut();
let mut count: c_uint = 0; let mut count: c_uint = 0;
@@ -71,9 +70,7 @@ impl OCREngine for AppleVisionOCR {
for box_data in boxes_slice { for box_data in boxes_slice {
// Convert C string to Rust String // Convert C string to Rust String
let text = if !box_data.text.is_null() { let text = if !box_data.text.is_null() {
CStr::from_ptr(box_data.text) CStr::from_ptr(box_data.text).to_string_lossy().into_owned()
.to_string_lossy()
.into_owned()
} else { } else {
String::new() String::new()
}; };
@@ -1,4 +1,4 @@
use crate::{ComputerController, types::*}; use crate::{types::*, ComputerController};
use anyhow::Result; use anyhow::Result;
use async_trait::async_trait; use async_trait::async_trait;
use tesseract::Tesseract; use tesseract::Tesseract;
@@ -62,7 +62,12 @@ impl ComputerController for LinuxController {
anyhow::bail!("Linux implementation not yet available") anyhow::bail!("Linux implementation not yet available")
} }
async fn take_screenshot(&self, _path: &str, _region: Option<Rect>, _window_id: Option<&str>) -> Result<()> { async fn take_screenshot(
&self,
_path: &str,
_region: Option<Rect>,
_window_id: Option<&str>,
) -> Result<()> {
// Enforce that window_id must be provided // Enforce that window_id must be provided
if _window_id.is_none() { if _window_id.is_none() {
anyhow::bail!("window_id is required. You must specify which window to capture (e.g., 'Firefox', 'Terminal', 'gedit'). Use list_windows to see available windows."); anyhow::bail!("window_id is required. You must specify which window to capture (e.g., 'Firefox', 'Terminal', 'gedit'). Use list_windows to see available windows.");
@@ -82,26 +87,31 @@ impl ComputerController for LinuxController {
.output(); .output();
if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() { if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() {
anyhow::bail!("Tesseract OCR is not installed on your system.\n\n\ anyhow::bail!(
"Tesseract OCR is not installed on your system.\n\n\
To install tesseract:\n \ To install tesseract:\n \
Ubuntu/Debian: sudo apt-get install tesseract-ocr\n \ Ubuntu/Debian: sudo apt-get install tesseract-ocr\n \
RHEL/CentOS: sudo yum install tesseract\n \ RHEL/CentOS: sudo yum install tesseract\n \
Arch Linux: sudo pacman -S tesseract\n\n\ Arch Linux: sudo pacman -S tesseract\n\n\
After installation, restart your terminal and try again."); After installation, restart your terminal and try again."
);
} }
// Initialize Tesseract // Initialize Tesseract
let tess = Tesseract::new(None, Some("eng")) let tess = Tesseract::new(None, Some("eng")).map_err(|e| {
.map_err(|e| { anyhow::anyhow!(
anyhow::anyhow!("Failed to initialize Tesseract: {}\n\n\ "Failed to initialize Tesseract: {}\n\n\
This usually means:\n1. Tesseract is not properly installed\n\ This usually means:\n1. Tesseract is not properly installed\n\
2. Language data files are missing\n\nTo fix:\n \ 2. Language data files are missing\n\nTo fix:\n \
Ubuntu/Debian: sudo apt-get install tesseract-ocr-eng\n \ Ubuntu/Debian: sudo apt-get install tesseract-ocr-eng\n \
RHEL/CentOS: sudo yum install tesseract-langpack-eng\n \ RHEL/CentOS: sudo yum install tesseract-langpack-eng\n \
Arch Linux: sudo pacman -S tesseract-data-eng", e) Arch Linux: sudo pacman -S tesseract-data-eng",
})?; e
)
})?;
let text = tess.set_image(_path) let text = tess
.set_image(_path)
.map_err(|e| anyhow::anyhow!("Failed to load image '{}': {}", _path, e))? .map_err(|e| anyhow::anyhow!("Failed to load image '{}': {}", _path, e))?
.get_text() .get_text()
.map_err(|e| anyhow::anyhow!("Failed to extract text from image: {}", e))?; .map_err(|e| anyhow::anyhow!("Failed to extract text from image: {}", e))?;
@@ -112,7 +122,12 @@ impl ComputerController for LinuxController {
Ok(OCRResult { Ok(OCRResult {
text, text,
confidence, confidence,
bounds: Rect { x: 0, y: 0, width: 0, height: 0 }, // Would need image dimensions bounds: Rect {
x: 0,
y: 0,
width: 0,
height: 0,
}, // Would need image dimensions
}) })
} }
@@ -123,12 +138,14 @@ impl ComputerController for LinuxController {
.output(); .output();
if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() { if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() {
anyhow::bail!("Tesseract OCR is not installed on your system.\n\n\ anyhow::bail!(
"Tesseract OCR is not installed on your system.\n\n\
To install tesseract:\n \ To install tesseract:\n \
Ubuntu/Debian: sudo apt-get install tesseract-ocr\n \ Ubuntu/Debian: sudo apt-get install tesseract-ocr\n \
RHEL/CentOS: sudo yum install tesseract\n \ RHEL/CentOS: sudo yum install tesseract\n \
Arch Linux: sudo pacman -S tesseract\n\n\ Arch Linux: sudo pacman -S tesseract\n\n\
After installation, restart your terminal and try again."); After installation, restart your terminal and try again."
);
} }
// Take full screen screenshot // Take full screen screenshot
@@ -136,17 +153,20 @@ impl ComputerController for LinuxController {
self.take_screenshot(&temp_path, None, None).await?; self.take_screenshot(&temp_path, None, None).await?;
// Use Tesseract to find text with bounding boxes // Use Tesseract to find text with bounding boxes
let tess = Tesseract::new(None, Some("eng")) let tess = Tesseract::new(None, Some("eng")).map_err(|e| {
.map_err(|e| { anyhow::anyhow!(
anyhow::anyhow!("Failed to initialize Tesseract: {}\n\n\ "Failed to initialize Tesseract: {}\n\n\
This usually means:\n1. Tesseract is not properly installed\n\ This usually means:\n1. Tesseract is not properly installed\n\
2. Language data files are missing\n\nTo fix:\n \ 2. Language data files are missing\n\nTo fix:\n \
Ubuntu/Debian: sudo apt-get install tesseract-ocr-eng\n \ Ubuntu/Debian: sudo apt-get install tesseract-ocr-eng\n \
RHEL/CentOS: sudo yum install tesseract-langpack-eng\n \ RHEL/CentOS: sudo yum install tesseract-langpack-eng\n \
Arch Linux: sudo pacman -S tesseract-data-eng", e) Arch Linux: sudo pacman -S tesseract-data-eng",
})?; e
)
})?;
let full_text = tess.set_image(temp_path.as_str()) let full_text = tess
.set_image(temp_path.as_str())
.map_err(|e| anyhow::anyhow!("Failed to load screenshot: {}", e))? .map_err(|e| anyhow::anyhow!("Failed to load screenshot: {}", e))?
.get_text() .get_text()
.map_err(|e| anyhow::anyhow!("Failed to extract text from screen: {}", e))?; .map_err(|e| anyhow::anyhow!("Failed to extract text from screen: {}", e))?;
@@ -157,7 +177,9 @@ impl ComputerController for LinuxController {
// Simple text search - full implementation would use get_component_images // Simple text search - full implementation would use get_component_images
// to get bounding boxes for each word // to get bounding boxes for each word
if full_text.contains(_text) { if full_text.contains(_text) {
tracing::warn!("Text found but precise coordinates not available in simplified implementation"); tracing::warn!(
"Text found but precise coordinates not available in simplified implementation"
);
Ok(Some(Point { x: 0, y: 0 })) Ok(Some(Point { x: 0, y: 0 }))
} else { } else {
Ok(None) Ok(None)
+219 -101
View File
@@ -1,13 +1,18 @@
use crate::{ComputerController, types::{Rect, TextLocation}}; use crate::ocr::{DefaultOCR, OCREngine};
use crate::ocr::{OCREngine, DefaultOCR}; use crate::{
use anyhow::{Result, Context}; types::{Rect, TextLocation},
ComputerController,
};
use anyhow::{Context, Result};
use async_trait::async_trait; use async_trait::async_trait;
use std::path::Path; use core_foundation::array::CFArray;
use core_graphics::window::{kCGWindowListOptionOnScreenOnly, kCGNullWindowID, CGWindowListCopyWindowInfo}; use core_foundation::base::{TCFType, ToVoid};
use core_foundation::dictionary::CFDictionary; use core_foundation::dictionary::CFDictionary;
use core_foundation::string::CFString; use core_foundation::string::CFString;
use core_foundation::base::{TCFType, ToVoid}; use core_graphics::window::{
use core_foundation::array::CFArray; kCGNullWindowID, kCGWindowListOptionOnScreenOnly, CGWindowListCopyWindowInfo,
};
use std::path::Path;
pub struct MacOSController { pub struct MacOSController {
ocr_engine: Box<dyn OCREngine>, ocr_engine: Box<dyn OCREngine>,
@@ -20,13 +25,21 @@ impl MacOSController {
let ocr = Box::new(DefaultOCR::new()?); let ocr = Box::new(DefaultOCR::new()?);
let ocr_name = ocr.name().to_string(); let ocr_name = ocr.name().to_string();
tracing::info!("Initialized macOS controller with OCR engine: {}", ocr_name); tracing::info!("Initialized macOS controller with OCR engine: {}", ocr_name);
Ok(Self { ocr_engine: ocr, ocr_name }) Ok(Self {
ocr_engine: ocr,
ocr_name,
})
} }
} }
#[async_trait] #[async_trait]
impl ComputerController for MacOSController { impl ComputerController for MacOSController {
async fn take_screenshot(&self, path: &str, region: Option<Rect>, window_id: Option<&str>) -> Result<()> { async fn take_screenshot(
&self,
path: &str,
region: Option<Rect>,
window_id: Option<&str>,
) -> Result<()> {
// Enforce that window_id must be provided // Enforce that window_id must be provided
if window_id.is_none() { if window_id.is_none() {
return Err(anyhow::anyhow!("window_id is required. You must specify which window to capture (e.g., 'Safari', 'Terminal', 'Google Chrome'). Use list_windows to see available windows.")); return Err(anyhow::anyhow!("window_id is required. You must specify which window to capture (e.g., 'Safari', 'Terminal', 'Google Chrome'). Use list_windows to see available windows."));
@@ -56,10 +69,8 @@ impl ComputerController for MacOSController {
// Get the window ID for the specified application // Get the window ID for the specified application
let cg_window_id = unsafe { let cg_window_id = unsafe {
let window_list = CGWindowListCopyWindowInfo( let window_list =
kCGWindowListOptionOnScreenOnly, CGWindowListCopyWindowInfo(kCGWindowListOptionOnScreenOnly, kCGNullWindowID);
kCGNullWindowID
);
let array = CFArray::<CFDictionary>::wrap_under_create_rule(window_list); let array = CFArray::<CFDictionary>::wrap_under_create_rule(window_list);
let count = array.len(); let count = array.len();
@@ -79,7 +90,11 @@ impl ComputerController for MacOSController {
continue; continue;
}; };
tracing::debug!("Checking window: owner='{}', looking for '{}'", owner, app_name); tracing::debug!(
"Checking window: owner='{}', looking for '{}'",
owner,
app_name
);
let owner_lower = owner.to_lowercase(); let owner_lower = owner.to_lowercase();
// Normalize by removing spaces for exact matching // Normalize by removing spaces for exact matching
@@ -88,18 +103,21 @@ impl ComputerController for MacOSController {
// ONLY accept exact matches (case-insensitive, with or without spaces) // ONLY accept exact matches (case-insensitive, with or without spaces)
// This prevents "Goose" from matching "GooseStudio" // This prevents "Goose" from matching "GooseStudio"
let is_match = owner_lower == app_name_lower || owner_normalized == app_name_normalized; let is_match =
owner_lower == app_name_lower || owner_normalized == app_name_normalized;
if is_match { if is_match {
// Get window ID // Get window ID
let window_id_key = CFString::from_static_string("kCGWindowNumber"); let window_id_key = CFString::from_static_string("kCGWindowNumber");
if let Some(value) = dict.find(window_id_key.to_void()) { if let Some(value) = dict.find(window_id_key.to_void()) {
let num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*value as *const _); let num: core_foundation::number::CFNumber =
TCFType::wrap_under_get_rule(*value as *const _);
if let Some(id) = num.to_i64() { if let Some(id) = num.to_i64() {
// Get window layer to filter out menu bar windows // Get window layer to filter out menu bar windows
let layer_key = CFString::from_static_string("kCGWindowLayer"); let layer_key = CFString::from_static_string("kCGWindowLayer");
let layer: i32 = if let Some(value) = dict.find(layer_key.to_void()) { let layer: i32 = if let Some(value) = dict.find(layer_key.to_void()) {
let num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*value as *const _); let num: core_foundation::number::CFNumber =
TCFType::wrap_under_get_rule(*value as *const _);
num.to_i32().unwrap_or(0) num.to_i32().unwrap_or(0)
} else { } else {
0 0
@@ -107,27 +125,31 @@ impl ComputerController for MacOSController {
// Get window bounds to verify it's a real window // Get window bounds to verify it's a real window
let bounds_key = CFString::from_static_string("kCGWindowBounds"); let bounds_key = CFString::from_static_string("kCGWindowBounds");
let has_real_bounds = if let Some(value) = dict.find(bounds_key.to_void()) { let has_real_bounds =
let bounds_dict: CFDictionary = TCFType::wrap_under_get_rule(*value as *const _); if let Some(value) = dict.find(bounds_key.to_void()) {
let width_key = CFString::from_static_string("Width"); let bounds_dict: CFDictionary =
let height_key = CFString::from_static_string("Height"); TCFType::wrap_under_get_rule(*value as *const _);
let width_key = CFString::from_static_string("Width");
let height_key = CFString::from_static_string("Height");
if let (Some(w_val), Some(h_val)) = ( if let (Some(w_val), Some(h_val)) = (
bounds_dict.find(width_key.to_void()), bounds_dict.find(width_key.to_void()),
bounds_dict.find(height_key.to_void()), bounds_dict.find(height_key.to_void()),
) { ) {
let w_num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*w_val as *const _); let w_num: core_foundation::number::CFNumber =
let h_num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*h_val as *const _); TCFType::wrap_under_get_rule(*w_val as *const _);
let width = w_num.to_f64().unwrap_or(0.0); let h_num: core_foundation::number::CFNumber =
let height = h_num.to_f64().unwrap_or(0.0); TCFType::wrap_under_get_rule(*h_val as *const _);
// Real windows should be at least 100x100 pixels let width = w_num.to_f64().unwrap_or(0.0);
width >= 100.0 && height >= 100.0 let height = h_num.to_f64().unwrap_or(0.0);
// Real windows should be at least 100x100 pixels
width >= 100.0 && height >= 100.0
} else {
false
}
} else { } else {
false false
} };
} else {
false
};
// Only accept windows that are: // Only accept windows that are:
// 1. At layer 0 (normal windows, not menu bar) // 1. At layer 0 (normal windows, not menu bar)
@@ -137,7 +159,13 @@ impl ComputerController for MacOSController {
found_window_id = Some((id as u32, owner.clone())); found_window_id = Some((id as u32, owner.clone()));
break; break;
} else { } else {
tracing::debug!("Skipping window ID {} for '{}': layer={}, has_real_bounds={}", id, owner, layer, has_real_bounds); tracing::debug!(
"Skipping window ID {} for '{}': layer={}, has_real_bounds={}",
id,
owner,
layer,
has_real_bounds
);
} }
} }
} }
@@ -150,7 +178,11 @@ impl ComputerController for MacOSController {
let (cg_window_id, matched_owner) = cg_window_id.ok_or_else(|| { let (cg_window_id, matched_owner) = cg_window_id.ok_or_else(|| {
anyhow::anyhow!("Could not find window for application '{}'. Use list_windows to see available windows.", app_name) anyhow::anyhow!("Could not find window for application '{}'. Use list_windows to see available windows.", app_name)
})?; })?;
tracing::info!("Taking screenshot of window ID {} for app '{}'", cg_window_id, matched_owner); tracing::info!(
"Taking screenshot of window ID {} for app '{}'",
cg_window_id,
matched_owner
);
// Use screencapture with the window ID for now // Use screencapture with the window ID for now
// TODO: Implement direct CGWindowListCreateImage approach with proper image saving // TODO: Implement direct CGWindowListCreateImage approach with proper image saving
@@ -161,7 +193,10 @@ impl ComputerController for MacOSController {
if let Some(region) = region { if let Some(region) = region {
cmd.arg("-R"); cmd.arg("-R");
cmd.arg(format!("{},{},{},{}", region.x, region.y, region.width, region.height)); cmd.arg(format!(
"{},{},{},{}",
region.x, region.y, region.width, region.height
));
} }
cmd.arg(&final_path); cmd.arg(&final_path);
@@ -170,7 +205,11 @@ impl ComputerController for MacOSController {
if !screenshot_result.status.success() { if !screenshot_result.status.success() {
let stderr = String::from_utf8_lossy(&screenshot_result.stderr); let stderr = String::from_utf8_lossy(&screenshot_result.stderr);
return Err(anyhow::anyhow!("screencapture failed for window {}: {}", cg_window_id, stderr)); return Err(anyhow::anyhow!(
"screencapture failed for window {}: {}",
cg_window_id,
stderr
));
} }
Ok(()) Ok(())
@@ -179,7 +218,8 @@ impl ComputerController for MacOSController {
async fn extract_text_from_screen(&self, region: Rect, window_id: &str) -> Result<String> { async fn extract_text_from_screen(&self, region: Rect, window_id: &str) -> Result<String> {
// Take screenshot of region first // Take screenshot of region first
let temp_path = format!("/tmp/g3_ocr_{}.png", uuid::Uuid::new_v4()); let temp_path = format!("/tmp/g3_ocr_{}.png", uuid::Uuid::new_v4());
self.take_screenshot(&temp_path, Some(region), Some(window_id)).await?; self.take_screenshot(&temp_path, Some(region), Some(window_id))
.await?;
// Extract text from the screenshot // Extract text from the screenshot
let result = self.extract_text_from_image(&temp_path).await?; let result = self.extract_text_from_image(&temp_path).await?;
@@ -193,7 +233,11 @@ impl ComputerController for MacOSController {
async fn extract_text_from_image(&self, path: &str) -> Result<String> { async fn extract_text_from_image(&self, path: &str) -> Result<String> {
// Extract all text and concatenate // Extract all text and concatenate
let locations = self.ocr_engine.extract_text_with_locations(path).await?; let locations = self.ocr_engine.extract_text_with_locations(path).await?;
Ok(locations.iter().map(|loc| loc.text.as_str()).collect::<Vec<_>>().join(" ")) Ok(locations
.iter()
.map(|loc| loc.text.as_str())
.collect::<Vec<_>>()
.join(" "))
} }
async fn extract_text_with_locations(&self, path: &str) -> Result<Vec<TextLocation>> { async fn extract_text_with_locations(&self, path: &str) -> Result<Vec<TextLocation>> {
@@ -201,11 +245,21 @@ impl ComputerController for MacOSController {
self.ocr_engine.extract_text_with_locations(path).await self.ocr_engine.extract_text_with_locations(path).await
} }
async fn find_text_in_app(&self, app_name: &str, search_text: &str) -> Result<Option<TextLocation>> { async fn find_text_in_app(
&self,
app_name: &str,
search_text: &str,
) -> Result<Option<TextLocation>> {
// Take screenshot of specific app window // Take screenshot of specific app window
let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string()); let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
let temp_path = format!("{}/tmp/g3_find_text_{}_{}.png", home, app_name, uuid::Uuid::new_v4()); let temp_path = format!(
self.take_screenshot(&temp_path, None, Some(app_name)).await?; "{}/tmp/g3_find_text_{}_{}.png",
home,
app_name,
uuid::Uuid::new_v4()
);
self.take_screenshot(&temp_path, None, Some(app_name))
.await?;
// Get screenshot dimensions before we delete it // Get screenshot dimensions before we delete it
let screenshot_dims = get_image_dimensions(&temp_path)?; let screenshot_dims = get_image_dimensions(&temp_path)?;
@@ -224,11 +278,8 @@ impl ComputerController for MacOSController {
for location in locations { for location in locations {
if location.text.to_lowercase().contains(&search_lower) { if location.text.to_lowercase().contains(&search_lower) {
// Transform coordinates from screenshot space to screen space // Transform coordinates from screenshot space to screen space
let transformed = transform_screenshot_to_screen_coords( let transformed =
location, transform_screenshot_to_screen_coords(location, window_bounds, screenshot_dims);
window_bounds,
screenshot_dims,
);
return Ok(Some(transformed)); return Ok(Some(transformed));
} }
} }
@@ -237,23 +288,22 @@ impl ComputerController for MacOSController {
} }
fn move_mouse(&self, x: i32, y: i32) -> Result<()> { fn move_mouse(&self, x: i32, y: i32) -> Result<()> {
use core_graphics::event::{ use core_graphics::event::{CGEvent, CGEventTapLocation, CGEventType, CGMouseButton};
CGEvent, CGEventTapLocation, CGEventType, CGMouseButton, use core_graphics::event_source::{CGEventSource, CGEventSourceStateID};
};
use core_graphics::event_source::{
CGEventSource, CGEventSourceStateID,
};
use core_graphics::geometry::CGPoint; use core_graphics::geometry::CGPoint;
let source = CGEventSource::new(CGEventSourceStateID::HIDSystemState) let source = CGEventSource::new(CGEventSourceStateID::HIDSystemState)
.ok().context("Failed to create event source")?; .ok()
.context("Failed to create event source")?;
let event = CGEvent::new_mouse_event( let event = CGEvent::new_mouse_event(
source, source,
CGEventType::MouseMoved, CGEventType::MouseMoved,
CGPoint::new(x as f64, y as f64), CGPoint::new(x as f64, y as f64),
CGMouseButton::Left, CGMouseButton::Left,
).ok().context("Failed to create mouse event")?; )
.ok()
.context("Failed to create mouse event")?;
event.post(CGEventTapLocation::HID); event.post(CGEventTapLocation::HID);
@@ -261,14 +311,10 @@ impl ComputerController for MacOSController {
} }
fn click_at(&self, x: i32, y: i32, _app_name: Option<&str>) -> Result<()> { fn click_at(&self, x: i32, y: i32, _app_name: Option<&str>) -> Result<()> {
use core_graphics::event::{
CGEvent, CGEventTapLocation, CGEventType, CGMouseButton,
};
use core_graphics::event_source::{
CGEventSource, CGEventSourceStateID,
};
use core_graphics::geometry::CGPoint;
use core_graphics::display::CGDisplay; use core_graphics::display::CGDisplay;
use core_graphics::event::{CGEvent, CGEventTapLocation, CGEventType, CGMouseButton};
use core_graphics::event_source::{CGEventSource, CGEventSourceStateID};
use core_graphics::geometry::CGPoint;
// IMPORTANT: Coordinates passed here are in NSScreen/CGWindowListCopyWindowInfo space // IMPORTANT: Coordinates passed here are in NSScreen/CGWindowListCopyWindowInfo space
// (Y=0 at BOTTOM, increases UPWARD) // (Y=0 at BOTTOM, increases UPWARD)
@@ -279,15 +325,22 @@ impl ComputerController for MacOSController {
let cgevent_x = x; let cgevent_x = x;
let cgevent_y = screen_height - y; let cgevent_y = screen_height - y;
tracing::debug!("click_at: NSScreen coords ({}, {}) -> CGEvent coords ({}, {}) [screen_height={}]", tracing::debug!(
x, y, cgevent_x, cgevent_y, screen_height); "click_at: NSScreen coords ({}, {}) -> CGEvent coords ({}, {}) [screen_height={}]",
x,
y,
cgevent_x,
cgevent_y,
screen_height
);
let (global_x, global_y) = (cgevent_x, cgevent_y); let (global_x, global_y) = (cgevent_x, cgevent_y);
let point = CGPoint::new(global_x as f64, global_y as f64); let point = CGPoint::new(global_x as f64, global_y as f64);
let source = CGEventSource::new(CGEventSourceStateID::HIDSystemState) let source = CGEventSource::new(CGEventSourceStateID::HIDSystemState)
.ok().context("Failed to create event source")?; .ok()
.context("Failed to create event source")?;
// Move mouse to position first // Move mouse to position first
let move_event = CGEvent::new_mouse_event( let move_event = CGEvent::new_mouse_event(
@@ -295,7 +348,9 @@ impl ComputerController for MacOSController {
CGEventType::MouseMoved, CGEventType::MouseMoved,
point, point,
CGMouseButton::Left, CGMouseButton::Left,
).ok().context("Failed to create mouse move event")?; )
.ok()
.context("Failed to create mouse move event")?;
move_event.post(CGEventTapLocation::HID); move_event.post(CGEventTapLocation::HID);
std::thread::sleep(std::time::Duration::from_millis(100)); std::thread::sleep(std::time::Duration::from_millis(100));
@@ -306,18 +361,18 @@ impl ComputerController for MacOSController {
CGEventType::LeftMouseDown, CGEventType::LeftMouseDown,
point, point,
CGMouseButton::Left, CGMouseButton::Left,
).ok().context("Failed to create mouse down event")?; )
.ok()
.context("Failed to create mouse down event")?;
mouse_down.post(CGEventTapLocation::HID); mouse_down.post(CGEventTapLocation::HID);
std::thread::sleep(std::time::Duration::from_millis(50)); std::thread::sleep(std::time::Duration::from_millis(50));
// Mouse up // Mouse up
let mouse_up = CGEvent::new_mouse_event( let mouse_up =
source, CGEvent::new_mouse_event(source, CGEventType::LeftMouseUp, point, CGMouseButton::Left)
CGEventType::LeftMouseUp, .ok()
point, .context("Failed to create mouse up event")?;
CGMouseButton::Left,
).ok().context("Failed to create mouse up event")?;
mouse_up.post(CGEventTapLocation::HID); mouse_up.post(CGEventTapLocation::HID);
Ok(()) Ok(())
@@ -328,10 +383,8 @@ impl MacOSController {
/// Get window bounds for an application (helper method) /// Get window bounds for an application (helper method)
fn get_window_bounds(&self, app_name: &str) -> Result<(i32, i32, i32, i32)> { fn get_window_bounds(&self, app_name: &str) -> Result<(i32, i32, i32, i32)> {
unsafe { unsafe {
let window_list = CGWindowListCopyWindowInfo( let window_list =
kCGWindowListOptionOnScreenOnly, CGWindowListCopyWindowInfo(kCGWindowListOptionOnScreenOnly, kCGNullWindowID);
kCGNullWindowID
);
let array = CFArray::<CFDictionary>::wrap_under_create_rule(window_list); let array = CFArray::<CFDictionary>::wrap_under_create_rule(window_list);
let count = array.len(); let count = array.len();
@@ -358,13 +411,15 @@ impl MacOSController {
// ONLY accept exact matches (case-insensitive, with or without spaces) // ONLY accept exact matches (case-insensitive, with or without spaces)
// This prevents "Goose" from matching "GooseStudio" // This prevents "Goose" from matching "GooseStudio"
let is_match = owner_lower == app_name_lower || owner_normalized == app_name_normalized; let is_match =
owner_lower == app_name_lower || owner_normalized == app_name_normalized;
if is_match { if is_match {
// Get window layer to filter out menu bar windows // Get window layer to filter out menu bar windows
let layer_key = CFString::from_static_string("kCGWindowLayer"); let layer_key = CFString::from_static_string("kCGWindowLayer");
let layer: i32 = if let Some(value) = dict.find(layer_key.to_void()) { let layer: i32 = if let Some(value) = dict.find(layer_key.to_void()) {
let num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*value as *const _); let num: core_foundation::number::CFNumber =
TCFType::wrap_under_get_rule(*value as *const _);
num.to_i32().unwrap_or(0) num.to_i32().unwrap_or(0)
} else { } else {
0 0
@@ -372,14 +427,19 @@ impl MacOSController {
// Skip menu bar windows (layer >= 20) // Skip menu bar windows (layer >= 20)
if layer >= 20 { if layer >= 20 {
tracing::debug!("Skipping window for '{}' at layer {} (menu bar)", owner, layer); tracing::debug!(
"Skipping window for '{}' at layer {} (menu bar)",
owner,
layer
);
continue; continue;
} }
// Get window bounds to verify it's a real window // Get window bounds to verify it's a real window
let bounds_key = CFString::from_static_string("kCGWindowBounds"); let bounds_key = CFString::from_static_string("kCGWindowBounds");
if let Some(value) = dict.find(bounds_key.to_void()) { if let Some(value) = dict.find(bounds_key.to_void()) {
let bounds_dict: CFDictionary = TCFType::wrap_under_get_rule(*value as *const _); let bounds_dict: CFDictionary =
TCFType::wrap_under_get_rule(*value as *const _);
let x_key = CFString::from_static_string("X"); let x_key = CFString::from_static_string("X");
let y_key = CFString::from_static_string("Y"); let y_key = CFString::from_static_string("Y");
@@ -392,10 +452,14 @@ impl MacOSController {
bounds_dict.find(width_key.to_void()), bounds_dict.find(width_key.to_void()),
bounds_dict.find(height_key.to_void()), bounds_dict.find(height_key.to_void()),
) { ) {
let x_num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*x_val as *const _); let x_num: core_foundation::number::CFNumber =
let y_num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*y_val as *const _); TCFType::wrap_under_get_rule(*x_val as *const _);
let w_num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*w_val as *const _); let y_num: core_foundation::number::CFNumber =
let h_num: core_foundation::number::CFNumber = TCFType::wrap_under_get_rule(*h_val as *const _); TCFType::wrap_under_get_rule(*y_val as *const _);
let w_num: core_foundation::number::CFNumber =
TCFType::wrap_under_get_rule(*w_val as *const _);
let h_num: core_foundation::number::CFNumber =
TCFType::wrap_under_get_rule(*h_val as *const _);
let x: i32 = x_num.to_i64().unwrap_or(0) as i32; let x: i32 = x_num.to_i64().unwrap_or(0) as i32;
let y: i32 = y_num.to_i64().unwrap_or(0) as i32; let y: i32 = y_num.to_i64().unwrap_or(0) as i32;
@@ -407,7 +471,12 @@ impl MacOSController {
tracing::info!("Found valid window bounds for '{}': x={}, y={}, w={}, h={} (layer={})", owner, x, y, w, h, layer); tracing::info!("Found valid window bounds for '{}': x={}, y={}, w={}, h={} (layer={})", owner, x, y, w, h, layer);
return Ok((x, y, w, h)); return Ok((x, y, w, h));
} else { } else {
tracing::debug!("Skipping window for '{}': too small ({}x{})", owner, w, h); tracing::debug!(
"Skipping window for '{}': too small ({}x{})",
owner,
w,
h
);
continue; continue;
} }
} else { } else {
@@ -418,7 +487,10 @@ impl MacOSController {
} }
} }
Err(anyhow::anyhow!("Could not find window bounds for '{}'", app_name)) Err(anyhow::anyhow!(
"Could not find window bounds for '{}'",
app_name
))
} }
} }
@@ -454,7 +526,7 @@ fn get_image_dimensions(path: &str) -> Result<(i32, i32)> {
fn transform_screenshot_to_screen_coords( fn transform_screenshot_to_screen_coords(
location: TextLocation, location: TextLocation,
window_bounds: (i32, i32, i32, i32), // (x, y, width, height) in screen space window_bounds: (i32, i32, i32, i32), // (x, y, width, height) in screen space
screenshot_dims: (i32, i32), // (width, height) in pixels screenshot_dims: (i32, i32), // (width, height) in pixels
) -> TextLocation { ) -> TextLocation {
let (win_x, win_y, win_width, win_height) = window_bounds; let (win_x, win_y, win_width, win_height) = window_bounds;
let (screenshot_width, screenshot_height) = screenshot_dims; let (screenshot_width, screenshot_height) = screenshot_dims;
@@ -464,8 +536,17 @@ fn transform_screenshot_to_screen_coords(
let scale_x = win_width as f64 / screenshot_width as f64; let scale_x = win_width as f64 / screenshot_width as f64;
let scale_y = win_height as f64 / screenshot_height as f64; let scale_y = win_height as f64 / screenshot_height as f64;
tracing::debug!("Transform: screenshot={}x{}, window={}x{} at ({},{}), scale=({:.2},{:.2})", tracing::debug!(
screenshot_width, screenshot_height, win_width, win_height, win_x, win_y, scale_x, scale_y); "Transform: screenshot={}x{}, window={}x{} at ({},{}), scale=({:.2},{:.2})",
screenshot_width,
screenshot_height,
win_width,
win_height,
win_x,
win_y,
scale_x,
scale_y
);
// Transform coordinates from image space to screen space // Transform coordinates from image space to screen space
// IMPORTANT: macOS screen coordinates have origin at BOTTOM-LEFT (Y increases upward) // IMPORTANT: macOS screen coordinates have origin at BOTTOM-LEFT (Y increases upward)
@@ -474,9 +555,18 @@ fn transform_screenshot_to_screen_coords(
// So we need to: (win_y + win_height) to get window TOP, then subtract screenshot_y // So we need to: (win_y + win_height) to get window TOP, then subtract screenshot_y
let window_top_y = win_y + win_height; let window_top_y = win_y + win_height;
tracing::debug!("[transform] Input location in image space: x={}, y={}, width={}, height={}", tracing::debug!(
location.x, location.y, location.width, location.height); "[transform] Input location in image space: x={}, y={}, width={}, height={}",
tracing::debug!("[transform] Scale factors: scale_x={:.4}, scale_y={:.4}", scale_x, scale_y); location.x,
location.y,
location.width,
location.height
);
tracing::debug!(
"[transform] Scale factors: scale_x={:.4}, scale_y={:.4}",
scale_x,
scale_y
);
let transformed_x = win_x + (location.x as f64 * scale_x) as i32; let transformed_x = win_x + (location.x as f64 * scale_x) as i32;
let transformed_y = window_top_y - (location.y as f64 * scale_y) as i32; let transformed_y = window_top_y - (location.y as f64 * scale_y) as i32;
@@ -484,13 +574,41 @@ fn transform_screenshot_to_screen_coords(
let transformed_height = (location.height as f64 * scale_y) as i32; let transformed_height = (location.height as f64 * scale_y) as i32;
tracing::debug!("[transform] Calculation details:"); tracing::debug!("[transform] Calculation details:");
tracing::debug!(" - transformed_x = {} + ({} * {:.4}) = {} + {:.2} = {}", win_x, location.x, scale_x, win_x, location.x as f64 * scale_x, transformed_x); tracing::debug!(
tracing::debug!(" - transformed_width = ({} * {:.4}) = {:.2} -> {}", location.width, scale_x, location.width as f64 * scale_x, transformed_width); " - transformed_x = {} + ({} * {:.4}) = {} + {:.2} = {}",
tracing::debug!(" - transformed_height = ({} * {:.4}) = {:.2} -> {}", location.height, scale_y, location.height as f64 * scale_y, transformed_height); win_x,
location.x,
scale_x,
win_x,
location.x as f64 * scale_x,
transformed_x
);
tracing::debug!(
" - transformed_width = ({} * {:.4}) = {:.2} -> {}",
location.width,
scale_x,
location.width as f64 * scale_x,
transformed_width
);
tracing::debug!(
" - transformed_height = ({} * {:.4}) = {:.2} -> {}",
location.height,
scale_y,
location.height as f64 * scale_y,
transformed_height
);
tracing::debug!("Transformed location: screenshot=({},{}) {}x{} -> screen=({},{}) {}x{}", tracing::debug!(
location.x, location.y, location.width, location.height, "Transformed location: screenshot=({},{}) {}x{} -> screen=({},{}) {}x{}",
transformed_x, transformed_y, transformed_width, transformed_height); location.x,
location.y,
location.width,
location.height,
transformed_x,
transformed_y,
transformed_width,
transformed_height
);
TextLocation { TextLocation {
text: location.text, text: location.text,
@@ -16,7 +16,7 @@ mod window_matching_tests {
("Visual Studio Code", "VisualStudioCode", true), ("Visual Studio Code", "VisualStudioCode", true),
("Google Chrome", "Google Chrome", true), ("Google Chrome", "Google Chrome", true),
("Safari", "Safari", true), ("Safari", "Safari", true),
("iTerm", "iTerm2", true), // fuzzy match ("iTerm", "iTerm2", true), // fuzzy match
("Code", "Visual Studio Code", true), // fuzzy match ("Code", "Visual Studio Code", true), // fuzzy match
]; ];
@@ -1,4 +1,4 @@
use crate::{ComputerController, types::*}; use crate::{types::*, ComputerController};
use anyhow::Result; use anyhow::Result;
use async_trait::async_trait; use async_trait::async_trait;
use tesseract::Tesseract; use tesseract::Tesseract;
@@ -61,7 +61,12 @@ impl ComputerController for WindowsController {
anyhow::bail!("Windows implementation not yet available") anyhow::bail!("Windows implementation not yet available")
} }
async fn take_screenshot(&self, _path: &str, _region: Option<Rect>, _window_id: Option<&str>) -> Result<()> { async fn take_screenshot(
&self,
_path: &str,
_region: Option<Rect>,
_window_id: Option<&str>,
) -> Result<()> {
// Enforce that window_id must be provided // Enforce that window_id must be provided
if _window_id.is_none() { if _window_id.is_none() {
anyhow::bail!("window_id is required. You must specify which window to capture (e.g., 'Chrome', 'Terminal', 'Notepad'). Use list_windows to see available windows."); anyhow::bail!("window_id is required. You must specify which window to capture (e.g., 'Chrome', 'Terminal', 'Notepad'). Use list_windows to see available windows.");
@@ -81,27 +86,32 @@ impl ComputerController for WindowsController {
.output(); .output();
if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() { if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() {
anyhow::bail!("Tesseract OCR is not installed on your system.\n\n\ anyhow::bail!(
"Tesseract OCR is not installed on your system.\n\n\
To install tesseract on Windows:\n \ To install tesseract on Windows:\n \
1. Download the installer from: https://github.com/UB-Mannheim/tesseract/wiki\n \ 1. Download the installer from: https://github.com/UB-Mannheim/tesseract/wiki\n \
2. Run the installer and follow the instructions\n \ 2. Run the installer and follow the instructions\n \
3. Add tesseract to your PATH environment variable\n \ 3. Add tesseract to your PATH environment variable\n \
4. Restart your terminal/command prompt\n\n\ 4. Restart your terminal/command prompt\n\n\
After installation, restart your terminal and try again."); After installation, restart your terminal and try again."
);
} }
// Initialize Tesseract // Initialize Tesseract
let tess = Tesseract::new(None, Some("eng")) let tess = Tesseract::new(None, Some("eng")).map_err(|e| {
.map_err(|e| { anyhow::anyhow!(
anyhow::anyhow!("Failed to initialize Tesseract: {}\n\n\ "Failed to initialize Tesseract: {}\n\n\
This usually means:\n1. Tesseract is not properly installed\n\ This usually means:\n1. Tesseract is not properly installed\n\
2. Language data files are missing\n\nTo fix:\n \ 2. Language data files are missing\n\nTo fix:\n \
1. Reinstall tesseract from https://github.com/UB-Mannheim/tesseract/wiki\n \ 1. Reinstall tesseract from https://github.com/UB-Mannheim/tesseract/wiki\n \
2. Make sure to select 'Additional language data' during installation\n \ 2. Make sure to select 'Additional language data' during installation\n \
3. Ensure tesseract is in your PATH", e) 3. Ensure tesseract is in your PATH",
})?; e
)
})?;
let text = tess.set_image(_path) let text = tess
.set_image(_path)
.map_err(|e| anyhow::anyhow!("Failed to load image '{}': {}", _path, e))? .map_err(|e| anyhow::anyhow!("Failed to load image '{}': {}", _path, e))?
.get_text() .get_text()
.map_err(|e| anyhow::anyhow!("Failed to extract text from image: {}", e))?; .map_err(|e| anyhow::anyhow!("Failed to extract text from image: {}", e))?;
@@ -112,7 +122,12 @@ impl ComputerController for WindowsController {
Ok(OCRResult { Ok(OCRResult {
text, text,
confidence, confidence,
bounds: Rect { x: 0, y: 0, width: 0, height: 0 }, // Would need image dimensions bounds: Rect {
x: 0,
y: 0,
width: 0,
height: 0,
}, // Would need image dimensions
}) })
} }
@@ -123,13 +138,15 @@ impl ComputerController for WindowsController {
.output(); .output();
if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() { if tesseract_check.is_err() || !tesseract_check.as_ref().unwrap().status.success() {
anyhow::bail!("Tesseract OCR is not installed on your system.\n\n\ anyhow::bail!(
"Tesseract OCR is not installed on your system.\n\n\
To install tesseract on Windows:\n \ To install tesseract on Windows:\n \
1. Download the installer from: https://github.com/UB-Mannheim/tesseract/wiki\n \ 1. Download the installer from: https://github.com/UB-Mannheim/tesseract/wiki\n \
2. Run the installer and follow the instructions\n \ 2. Run the installer and follow the instructions\n \
3. Add tesseract to your PATH environment variable\n \ 3. Add tesseract to your PATH environment variable\n \
4. Restart your terminal/command prompt\n\n\ 4. Restart your terminal/command prompt\n\n\
After installation, restart your terminal and try again."); After installation, restart your terminal and try again."
);
} }
// Take full screen screenshot // Take full screen screenshot
@@ -137,17 +154,20 @@ impl ComputerController for WindowsController {
self.take_screenshot(&temp_path, None, None).await?; self.take_screenshot(&temp_path, None, None).await?;
// Use Tesseract to find text with bounding boxes // Use Tesseract to find text with bounding boxes
let tess = Tesseract::new(None, Some("eng")) let tess = Tesseract::new(None, Some("eng")).map_err(|e| {
.map_err(|e| { anyhow::anyhow!(
anyhow::anyhow!("Failed to initialize Tesseract: {}\n\n\ "Failed to initialize Tesseract: {}\n\n\
This usually means:\n1. Tesseract is not properly installed\n\ This usually means:\n1. Tesseract is not properly installed\n\
2. Language data files are missing\n\nTo fix:\n \ 2. Language data files are missing\n\nTo fix:\n \
1. Reinstall tesseract from https://github.com/UB-Mannheim/tesseract/wiki\n \ 1. Reinstall tesseract from https://github.com/UB-Mannheim/tesseract/wiki\n \
2. Make sure to select 'Additional language data' during installation\n \ 2. Make sure to select 'Additional language data' during installation\n \
3. Ensure tesseract is in your PATH", e) 3. Ensure tesseract is in your PATH",
})?; e
)
})?;
let full_text = tess.set_image(temp_path.as_str()) let full_text = tess
.set_image(temp_path.as_str())
.map_err(|e| anyhow::anyhow!("Failed to load screenshot: {}", e))? .map_err(|e| anyhow::anyhow!("Failed to load screenshot: {}", e))?
.get_text() .get_text()
.map_err(|e| anyhow::anyhow!("Failed to extract text from screen: {}", e))?; .map_err(|e| anyhow::anyhow!("Failed to extract text from screen: {}", e))?;
@@ -158,7 +178,9 @@ impl ComputerController for WindowsController {
// Simple text search - full implementation would use get_component_images // Simple text search - full implementation would use get_component_images
// to get bounding boxes for each word // to get bounding boxes for each word
if full_text.contains(_text) { if full_text.contains(_text) {
tracing::warn!("Text found but precise coordinates not available in simplified implementation"); tracing::warn!(
"Text found but precise coordinates not available in simplified implementation"
);
Ok(Some(Point { x: 0, y: 0 })) Ok(Some(Point { x: 0, y: 0 }))
} else { } else {
Ok(None) Ok(None)
@@ -105,7 +105,13 @@ impl WebElement {
/// Find multiple child elements by CSS selector /// Find multiple child elements by CSS selector
pub async fn find_elements(&mut self, selector: &str) -> Result<Vec<WebElement>> { pub async fn find_elements(&mut self, selector: &str) -> Result<Vec<WebElement>> {
let elems = self.inner.find_all(fantoccini::Locator::Css(selector)).await?; let elems = self
Ok(elems.into_iter().map(|inner| WebElement { inner }).collect()) .inner
.find_all(fantoccini::Locator::Css(selector))
.await?;
Ok(elems
.into_iter()
.map(|inner| WebElement { inner })
.collect())
} }
} }
@@ -29,7 +29,10 @@ impl SafariDriver {
let url = format!("http://localhost:{}", port); let url = format!("http://localhost:{}", port);
let mut caps = serde_json::Map::new(); let mut caps = serde_json::Map::new();
caps.insert("browserName".to_string(), Value::String("safari".to_string())); caps.insert(
"browserName".to_string(),
Value::String("safari".to_string()),
);
let client = ClientBuilder::native() let client = ClientBuilder::native()
.capabilities(caps) .capabilities(caps)
@@ -61,9 +64,7 @@ impl SafariDriver {
/// Get all window handles /// Get all window handles
pub async fn window_handles(&mut self) -> Result<Vec<String>> { pub async fn window_handles(&mut self) -> Result<Vec<String>> {
let handles = self.client.windows().await?; let handles = self.client.windows().await?;
Ok(handles.into_iter() Ok(handles.into_iter().map(|h| h.into()).collect())
.map(|h| h.into())
.collect())
} }
/// Switch to a window by handle /// Switch to a window by handle
@@ -109,7 +110,11 @@ impl SafariDriver {
} }
/// Wait for an element to appear (with timeout) /// Wait for an element to appear (with timeout)
pub async fn wait_for_element(&mut self, selector: &str, timeout: Duration) -> Result<WebElement> { pub async fn wait_for_element(
&mut self,
selector: &str,
timeout: Duration,
) -> Result<WebElement> {
let start = std::time::Instant::now(); let start = std::time::Instant::now();
let poll_interval = Duration::from_millis(100); let poll_interval = Duration::from_millis(100);
@@ -127,7 +132,11 @@ impl SafariDriver {
} }
/// Wait for an element to be visible (with timeout) /// Wait for an element to be visible (with timeout)
pub async fn wait_for_visible(&mut self, selector: &str, timeout: Duration) -> Result<WebElement> { pub async fn wait_for_visible(
&mut self,
selector: &str,
timeout: Duration,
) -> Result<WebElement> {
let start = std::time::Instant::now(); let start = std::time::Instant::now();
let poll_interval = Duration::from_millis(100); let poll_interval = Duration::from_millis(100);
@@ -163,14 +172,26 @@ impl WebDriverController for SafariDriver {
} }
async fn find_element(&mut self, selector: &str) -> Result<WebElement> { async fn find_element(&mut self, selector: &str) -> Result<WebElement> {
let elem = self.client.find(fantoccini::Locator::Css(selector)).await let elem = self
.context(format!("Failed to find element with selector: {}", selector))?; .client
.find(fantoccini::Locator::Css(selector))
.await
.context(format!(
"Failed to find element with selector: {}",
selector
))?;
Ok(WebElement { inner: elem }) Ok(WebElement { inner: elem })
} }
async fn find_elements(&mut self, selector: &str) -> Result<Vec<WebElement>> { async fn find_elements(&mut self, selector: &str) -> Result<Vec<WebElement>> {
let elems = self.client.find_all(fantoccini::Locator::Css(selector)).await?; let elems = self
Ok(elems.into_iter().map(|inner| WebElement { inner }).collect()) .client
.find_all(fantoccini::Locator::Css(selector))
.await?;
Ok(elems
.into_iter()
.map(|inner| WebElement { inner })
.collect())
} }
async fn execute_script(&mut self, script: &str, args: Vec<Value>) -> Result<Value> { async fn execute_script(&mut self, script: &str, args: Vec<Value>) -> Result<Value> {
@@ -194,8 +215,7 @@ impl WebDriverController for SafariDriver {
.context("Failed to create parent directories for screenshot")?; .context("Failed to create parent directories for screenshot")?;
} }
std::fs::write(path_str, screenshot_data) std::fs::write(path_str, screenshot_data).context("Failed to write screenshot to file")?;
.context("Failed to write screenshot to file")?;
Ok(()) Ok(())
} }
@@ -7,11 +7,17 @@ async fn test_screenshot() {
// Test that screenshot without window_id fails with appropriate error // Test that screenshot without window_id fails with appropriate error
let path = "/tmp/test_screenshot.png"; let path = "/tmp/test_screenshot.png";
let result = controller.take_screenshot(path, None, None).await; let result = controller.take_screenshot(path, None, None).await;
assert!(result.is_err(), "Expected error when window_id is not provided"); assert!(
result.is_err(),
"Expected error when window_id is not provided"
);
let error_msg = result.unwrap_err().to_string(); let error_msg = result.unwrap_err().to_string();
assert!(error_msg.contains("window_id is required"), assert!(
"Expected error message about window_id being required, got: {}", error_msg); error_msg.contains("window_id is required"),
"Expected error message about window_id being required, got: {}",
error_msg
);
} }
#[tokio::test] #[tokio::test]
+21 -25
View File
@@ -1,5 +1,5 @@
use serde::{Deserialize, Serialize};
use anyhow::Result; use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::path::Path; use std::path::Path;
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -21,7 +21,7 @@ pub struct ProvidersConfig {
pub databricks: Option<DatabricksConfig>, pub databricks: Option<DatabricksConfig>,
pub embedded: Option<EmbeddedConfig>, pub embedded: Option<EmbeddedConfig>,
pub default_provider: String, pub default_provider: String,
pub coach: Option<String>, // Provider to use for coach in autonomous mode pub coach: Option<String>, // Provider to use for coach in autonomous mode
pub player: Option<String>, // Provider to use for player in autonomous mode pub player: Option<String>, // Provider to use for player in autonomous mode
} }
@@ -103,9 +103,7 @@ pub struct MacAxConfig {
impl Default for MacAxConfig { impl Default for MacAxConfig {
fn default() -> Self { fn default() -> Self {
Self { Self { enabled: false }
enabled: false,
}
} }
} }
@@ -173,11 +171,7 @@ impl Config {
Path::new(path).exists() Path::new(path).exists()
} else { } else {
// Check default locations // Check default locations
let default_paths = [ let default_paths = ["./g3.toml", "~/.config/g3/config.toml", "~/.g3.toml"];
"./g3.toml",
"~/.config/g3/config.toml",
"~/.g3.toml",
];
default_paths.iter().any(|path| { default_paths.iter().any(|path| {
let expanded_path = shellexpand::tilde(path); let expanded_path = shellexpand::tilde(path);
@@ -205,7 +199,10 @@ impl Config {
if let Err(e) = databricks_config.save(config_file.to_str().unwrap()) { if let Err(e) = databricks_config.save(config_file.to_str().unwrap()) {
eprintln!("Warning: Could not save default config: {}", e); eprintln!("Warning: Could not save default config: {}", e);
} else { } else {
println!("Created default Databricks configuration at: {}", config_file.display()); println!(
"Created default Databricks configuration at: {}",
config_file.display()
);
} }
return Ok(databricks_config); return Ok(databricks_config);
@@ -224,11 +221,7 @@ impl Config {
} }
} else { } else {
// Try to load from default locations // Try to load from default locations
let default_paths = [ let default_paths = ["./g3.toml", "~/.config/g3/config.toml", "~/.g3.toml"];
"./g3.toml",
"~/.config/g3/config.toml",
"~/.g3.toml",
];
for path in &default_paths { for path in &default_paths {
let expanded_path = shellexpand::tilde(path); let expanded_path = shellexpand::tilde(path);
@@ -240,10 +233,7 @@ impl Config {
} }
// Override with environment variables // Override with environment variables
settings = settings.add_source( settings = settings.add_source(config::Environment::with_prefix("G3").separator("_"));
config::Environment::with_prefix("G3")
.separator("_")
);
let config = settings.build()?.try_deserialize()?; let config = settings.build()?.try_deserialize()?;
Ok(config) Ok(config)
@@ -260,7 +250,7 @@ impl Config {
embedded: Some(EmbeddedConfig { embedded: Some(EmbeddedConfig {
model_path: "~/.cache/g3/models/qwen2.5-7b-instruct-q3_k_m.gguf".to_string(), model_path: "~/.cache/g3/models/qwen2.5-7b-instruct-q3_k_m.gguf".to_string(),
model_type: "qwen".to_string(), model_type: "qwen".to_string(),
context_length: Some(32768), // Qwen2.5 supports 32k context context_length: Some(32768), // Qwen2.5 supports 32k context
max_tokens: Some(2048), max_tokens: Some(2048),
temperature: Some(0.1), temperature: Some(0.1),
gpu_layers: Some(32), gpu_layers: Some(32),
@@ -345,8 +335,12 @@ impl Config {
)); ));
} }
} }
_ => return Err(anyhow::anyhow!("Unknown provider: {}", _ => {
config.providers.default_provider)), return Err(anyhow::anyhow!(
"Unknown provider: {}",
config.providers.default_provider
))
}
} }
} }
@@ -355,14 +349,16 @@ impl Config {
/// Get the provider to use for coach mode in autonomous execution /// Get the provider to use for coach mode in autonomous execution
pub fn get_coach_provider(&self) -> &str { pub fn get_coach_provider(&self) -> &str {
self.providers.coach self.providers
.coach
.as_deref() .as_deref()
.unwrap_or(&self.providers.default_provider) .unwrap_or(&self.providers.default_provider)
} }
/// Get the provider to use for player mode in autonomous execution /// Get the provider to use for player mode in autonomous execution
pub fn get_player_provider(&self) -> &str { pub fn get_player_provider(&self) -> &str {
self.providers.player self.providers
.player
.as_deref() .as_deref()
.unwrap_or(&self.providers.default_provider) .unwrap_or(&self.providers.default_provider)
} }
@@ -1,6 +1,6 @@
#[cfg(test)] #[cfg(test)]
mod test_multiple_tool_calls { mod test_multiple_tool_calls {
use g3_config::{Config, AgentConfig}; use g3_config::{AgentConfig, Config};
#[test] #[test]
fn test_config_has_multiple_tool_calls_field() { fn test_config_has_multiple_tool_calls_field() {
+1 -1
View File
@@ -1,4 +1,4 @@
use sysinfo::{System, Pid}; use sysinfo::{Pid, System};
fn main() { fn main() {
let mut sys = System::new_all(); let mut sys = System::new_all();
+4 -2
View File
@@ -8,8 +8,10 @@ fn main() {
Ok(instances) => { Ok(instances) => {
println!("Found {} instances:", instances.len()); println!("Found {} instances:", instances.len());
for instance in instances { for instance in instances {
println!(" - PID: {}, Workspace: {:?}, Type: {:?}", println!(
instance.pid, instance.workspace, instance.instance_type); " - PID: {}, Workspace: {:?}, Type: {:?}",
instance.pid, instance.workspace, instance.instance_type
);
} }
} }
Err(e) => { Err(e) => {
+1 -1
View File
@@ -1,4 +1,4 @@
use sysinfo::{System, Pid}; use sysinfo::{Pid, System};
fn main() { fn main() {
let mut sys = System::new_all(); let mut sys = System::new_all();
+43 -28
View File
@@ -50,21 +50,24 @@ pub async fn restart_instance(
let mut controller = controller.lock().await; let mut controller = controller.lock().await;
// Get stored launch params // Get stored launch params
let params = controller.get_launch_params(pid) let params = controller
.get_launch_params(pid)
.ok_or(StatusCode::NOT_FOUND)?; .ok_or(StatusCode::NOT_FOUND)?;
// Launch new instance with same parameters // Launch new instance with same parameters
let new_pid = controller.launch_g3( let new_pid = controller
params.workspace.to_str().unwrap(), .launch_g3(
&params.provider, params.workspace.to_str().unwrap(),
&params.model, &params.provider,
&params.prompt, &params.model,
params.autonomous, &params.prompt,
params.g3_binary_path.as_deref(), params.autonomous,
).map_err(|e| { params.g3_binary_path.as_deref(),
error!("Failed to restart instance: {}", e); )
StatusCode::INTERNAL_SERVER_ERROR .map_err(|e| {
})?; error!("Failed to restart instance: {}", e);
StatusCode::INTERNAL_SERVER_ERROR
})?;
let new_id = format!("{}_{}", new_pid, chrono::Utc::now().timestamp()); let new_id = format!("{}_{}", new_pid, chrono::Utc::now().timestamp());
@@ -94,10 +97,13 @@ pub async fn launch_instance(
// Check if file exists // Check if file exists
if !path.exists() { if !path.exists() {
error!("G3 binary not found: {}", binary_path); error!("G3 binary not found: {}", binary_path);
return Err((StatusCode::BAD_REQUEST, Json(serde_json::json!({ return Err((
"error": "G3 binary not found", StatusCode::BAD_REQUEST,
"message": format!("The specified g3 binary does not exist: {}", binary_path) Json(serde_json::json!({
})))); "error": "G3 binary not found",
"message": format!("The specified g3 binary does not exist: {}", binary_path)
})),
));
} }
// Check if file is executable (Unix only) // Check if file is executable (Unix only)
@@ -107,20 +113,26 @@ pub async fn launch_instance(
if let Ok(metadata) = std::fs::metadata(path) { if let Ok(metadata) = std::fs::metadata(path) {
if metadata.permissions().mode() & 0o111 == 0 { if metadata.permissions().mode() & 0o111 == 0 {
error!("G3 binary is not executable: {}", binary_path); error!("G3 binary is not executable: {}", binary_path);
return Err((StatusCode::BAD_REQUEST, Json(serde_json::json!({ return Err((
"error": "G3 binary is not executable", StatusCode::BAD_REQUEST,
"message": format!("The specified g3 binary is not executable: {}", binary_path) Json(serde_json::json!({
})))); "error": "G3 binary is not executable",
"message": format!("The specified g3 binary is not executable: {}", binary_path)
})),
));
} }
} }
} }
} }
let workspace = request.workspace.to_str().ok_or_else(|| { let workspace = request.workspace.to_str().ok_or_else(|| {
(StatusCode::BAD_REQUEST, Json(serde_json::json!({ (
"error": "Invalid workspace path", StatusCode::BAD_REQUEST,
"message": "The workspace path contains invalid characters" Json(serde_json::json!({
}))) "error": "Invalid workspace path",
"message": "The workspace path contains invalid characters"
})),
)
})?; })?;
let autonomous = request.mode == LaunchMode::Ensemble; let autonomous = request.mode == LaunchMode::Ensemble;
let g3_binary_path = request.g3_binary_path.as_deref(); let g3_binary_path = request.g3_binary_path.as_deref();
@@ -145,10 +157,13 @@ pub async fn launch_instance(
} }
Err(e) => { Err(e) => {
error!("Failed to launch g3 instance: {}", e); error!("Failed to launch g3 instance: {}", e);
Err((StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ Err((
"error": "Failed to launch instance", StatusCode::INTERNAL_SERVER_ERROR,
"message": format!("Error: {}", e) Json(serde_json::json!({
})))) "error": "Failed to launch instance",
"message": format!("Error: {}", e)
})),
))
} }
} }
} }
+26 -17
View File
@@ -1,7 +1,11 @@
use crate::logs::{LogParser, StatsAggregator}; use crate::logs::{LogParser, StatsAggregator};
use crate::models::*; use crate::models::*;
use crate::process::ProcessDetector; use crate::process::ProcessDetector;
use axum::{extract::{Query, State}, http::StatusCode, Json}; use axum::{
extract::{Query, State},
http::StatusCode,
Json,
};
use serde::Deserialize; use serde::Deserialize;
use std::sync::Arc; use std::sync::Arc;
use tokio::sync::Mutex; use tokio::sync::Mutex;
@@ -69,7 +73,10 @@ fn get_instance_detail(instance: &Instance) -> anyhow::Result<InstanceDetail> {
let log_entries = match LogParser::parse_logs(&instance.workspace) { let log_entries = match LogParser::parse_logs(&instance.workspace) {
Ok(entries) => entries, Ok(entries) => entries,
Err(e) => { Err(e) => {
warn!("Failed to parse logs for instance {}: {}. Instance may be newly started.", instance.id, e); warn!(
"Failed to parse logs for instance {}: {}. Instance may be newly started.",
instance.id, e
);
Vec::new() Vec::new()
} }
}; };
@@ -85,7 +92,10 @@ fn get_instance_detail(instance: &Instance) -> anyhow::Result<InstanceDetail> {
let git_status = match get_git_status(&instance.workspace) { let git_status = match get_git_status(&instance.workspace) {
Some(status) => Some(status), Some(status) => Some(status),
None => { None => {
debug!("No git status available for workspace: {:?}", instance.workspace); debug!(
"No git status available for workspace: {:?}",
instance.workspace
);
None None
} }
}; };
@@ -177,16 +187,10 @@ fn read_file_snippet(workspace: &std::path::Path, filename: &str) -> Option<Stri
return None; return None;
} }
fs::read_to_string(&path) fs::read_to_string(&path).ok().map(|content| {
.ok() // Return first 10 lines
.map(|content| { content.lines().take(10).collect::<Vec<_>>().join("\n")
// Return first 10 lines })
content
.lines()
.take(10)
.collect::<Vec<_>>()
.join("\n")
})
} }
#[derive(Deserialize)] #[derive(Deserialize)]
@@ -202,8 +206,13 @@ pub async fn get_file_content(
let mut detector = detector.lock().await; let mut detector = detector.lock().await;
// Find the instance // Find the instance
let instances = detector.detect_instances().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let instances = detector
let instance = instances.iter().find(|i| i.id == id).ok_or(StatusCode::NOT_FOUND)?; .detect_instances()
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let instance = instances
.iter()
.find(|i| i.id == id)
.ok_or(StatusCode::NOT_FOUND)?;
// Read the full file // Read the full file
let file_path = instance.workspace.join(&query.name); let file_path = instance.workspace.join(&query.name);
@@ -211,8 +220,8 @@ pub async fn get_file_content(
return Err(StatusCode::NOT_FOUND); return Err(StatusCode::NOT_FOUND);
} }
let content = std::fs::read_to_string(&file_path) let content =
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; std::fs::read_to_string(&file_path).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
Ok(Json(serde_json::json!({ Ok(Json(serde_json::json!({
"name": query.name, "name": query.name,
+1 -1
View File
@@ -1,4 +1,4 @@
pub mod instances;
pub mod control; pub mod control;
pub mod instances;
pub mod logs; pub mod logs;
pub mod state; pub mod state;
+9 -9
View File
@@ -1,8 +1,8 @@
use crate::launch::ConsoleState; use crate::launch::ConsoleState;
use axum::{http::StatusCode, Json}; use axum::{http::StatusCode, Json};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::os::unix::fs::PermissionsExt; use std::os::unix::fs::PermissionsExt;
use std::path::PathBuf;
use tracing::{error, info}; use tracing::{error, info};
pub async fn get_state() -> Result<Json<ConsoleState>, StatusCode> { pub async fn get_state() -> Result<Json<ConsoleState>, StatusCode> {
@@ -59,12 +59,14 @@ pub async fn browse_filesystem(
std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")) std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
}; };
let current_path = path.canonicalize() let current_path = path
.canonicalize()
.map_err(|_| StatusCode::BAD_REQUEST)? .map_err(|_| StatusCode::BAD_REQUEST)?
.to_string_lossy() .to_string_lossy()
.to_string(); .to_string();
let parent_path = path.parent() let parent_path = path
.parent()
.and_then(|p| p.to_str()) .and_then(|p| p.to_str())
.map(|s| s.to_string()); .map(|s| s.to_string());
@@ -83,12 +85,10 @@ pub async fn browse_filesystem(
} }
} }
entries.sort_by(|a, b| { entries.sort_by(|a, b| match (a.is_dir, b.is_dir) {
match (a.is_dir, b.is_dir) { (true, false) => std::cmp::Ordering::Less,
(true, false) => std::cmp::Ordering::Less, (false, true) => std::cmp::Ordering::Greater,
(false, true) => std::cmp::Ordering::Greater, _ => a.name.cmp(&b.name),
_ => a.name.cmp(&b.name),
}
}); });
Ok(Json(BrowseResponse { Ok(Json(BrowseResponse {
+1 -3
View File
@@ -59,8 +59,6 @@ impl ConsoleState {
fn config_path() -> PathBuf { fn config_path() -> PathBuf {
// Use explicit ~/.config/g3/console.json path as per requirements // Use explicit ~/.config/g3/console.json path as per requirements
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from(".")); let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
home.join(".config") home.join(".config").join("g3").join("console.json")
.join("g3")
.join("console.json")
} }
} }
+1 -1
View File
@@ -1,5 +1,5 @@
pub mod api; pub mod api;
pub mod launch;
pub mod logs; pub mod logs;
pub mod models; pub mod models;
pub mod process; pub mod process;
pub mod launch;
+25 -15
View File
@@ -55,17 +55,21 @@ impl LogParser {
if let Some(messages) = json.get("messages").and_then(|m| m.as_array()) { if let Some(messages) = json.get("messages").and_then(|m| m.as_array()) {
for msg in messages { for msg in messages {
entries.push(LogEntry { entries.push(LogEntry {
timestamp: msg.get("timestamp") timestamp: msg
.get("timestamp")
.and_then(|t| t.as_str()) .and_then(|t| t.as_str())
.and_then(|s| DateTime::parse_from_rfc3339(s).ok()) .and_then(|s| DateTime::parse_from_rfc3339(s).ok())
.map(|dt| dt.with_timezone(&Utc)), .map(|dt| dt.with_timezone(&Utc)),
role: msg.get("role") role: msg
.get("role")
.and_then(|r| r.as_str()) .and_then(|r| r.as_str())
.map(String::from), .map(String::from),
content: msg.get("content") content: msg
.get("content")
.and_then(|c| c.as_str()) .and_then(|c| c.as_str())
.map(String::from), .map(String::from),
tool_calls: msg.get("tool_calls") tool_calls: msg
.get("tool_calls")
.and_then(|tc| tc.as_array()) .and_then(|tc| tc.as_array())
.map(|arr| arr.clone()), .map(|arr| arr.clone()),
raw: msg.clone(), raw: msg.clone(),
@@ -78,13 +82,11 @@ impl LogParser {
} }
// Sort by timestamp // Sort by timestamp
entries.sort_by(|a, b| { entries.sort_by(|a, b| match (&a.timestamp, &b.timestamp) {
match (&a.timestamp, &b.timestamp) { (Some(t1), Some(t2)) => t1.cmp(t2),
(Some(t1), Some(t2)) => t1.cmp(t2), (Some(_), None) => std::cmp::Ordering::Less,
(Some(_), None) => std::cmp::Ordering::Less, (None, Some(_)) => std::cmp::Ordering::Greater,
(None, Some(_)) => std::cmp::Ordering::Greater, (None, None) => std::cmp::Ordering::Equal,
(None, None) => std::cmp::Ordering::Equal,
}
}); });
Ok(entries) Ok(entries)
@@ -117,10 +119,12 @@ impl LogParser {
if let Some(name) = call.get("name").and_then(|n| n.as_str()) { if let Some(name) = call.get("name").and_then(|n| n.as_str()) {
tool_calls.push(ToolCall { tool_calls.push(ToolCall {
name: name.to_string(), name: name.to_string(),
parameters: call.get("parameters") parameters: call
.get("parameters")
.cloned() .cloned()
.unwrap_or(Value::Object(serde_json::Map::new())), .unwrap_or(Value::Object(serde_json::Map::new())),
result: call.get("result") result: call
.get("result")
.and_then(|r| r.as_str()) .and_then(|r| r.as_str())
.map(String::from), .map(String::from),
timestamp: entry.timestamp, timestamp: entry.timestamp,
@@ -193,7 +197,9 @@ impl StatsAggregator {
entries entries
.iter() .iter()
.filter_map(|entry| { .filter_map(|entry| {
entry.raw.get("usage") entry
.raw
.get("usage")
.and_then(|u| u.get("total_tokens")) .and_then(|u| u.get("total_tokens"))
.and_then(|t| t.as_u64()) .and_then(|t| t.as_u64())
}) })
@@ -213,7 +219,11 @@ impl StatsAggregator {
.iter() .iter()
.filter(|entry| { .filter(|entry| {
entry.raw.get("error").is_some() entry.raw.get("error").is_some()
|| entry.content.as_ref().map(|c| c.to_lowercase().contains("error")).unwrap_or(false) || entry
.content
.as_ref()
.map(|c| c.to_lowercase().contains("error"))
.unwrap_or(false)
}) })
.count() as u64 .count() as u64
} }
+4 -6
View File
@@ -1,11 +1,11 @@
use g3_console::api; use g3_console::api;
use g3_console::process;
use g3_console::launch; use g3_console::launch;
use g3_console::process;
use api::control::{kill_instance, launch_instance, restart_instance}; use api::control::{kill_instance, launch_instance, restart_instance};
use api::instances::{get_instance, get_file_content, list_instances}; use api::instances::{get_file_content, get_instance, list_instances};
use api::logs::get_instance_logs; use api::logs::get_instance_logs;
use api::state::{get_state, save_state, browse_filesystem}; use api::state::{browse_filesystem, get_state, save_state};
use axum::{ use axum::{
routing::{get, post}, routing::{get, post},
Router, Router,
@@ -39,9 +39,7 @@ struct Args {
#[tokio::main] #[tokio::main]
async fn main() -> anyhow::Result<()> { async fn main() -> anyhow::Result<()> {
// Initialize tracing // Initialize tracing
tracing_subscriber::fmt() tracing_subscriber::fmt().with_max_level(Level::INFO).init();
.with_max_level(Level::INFO)
.init();
let args = Args::parse(); let args = Args::parse();
+1 -1
View File
@@ -1,6 +1,6 @@
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::path::PathBuf; use std::path::PathBuf;
use chrono::{DateTime, Utc};
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Instance { pub struct Instance {
+1 -1
View File
@@ -1,5 +1,5 @@
use serde::{Deserialize, Serialize};
use chrono::{DateTime, Utc}; use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage { pub struct ChatMessage {
+19 -12
View File
@@ -1,12 +1,12 @@
use anyhow::{anyhow, Context, Result};
use std::process::{Command, Stdio};
use std::os::unix::process::CommandExt;
use std::collections::HashMap;
use std::sync::Mutex;
use std::path::PathBuf;
use sysinfo::{Pid, Signal, System, Process};
use tracing::{debug, info};
use crate::models::LaunchParams; use crate::models::LaunchParams;
use anyhow::{anyhow, Context, Result};
use std::collections::HashMap;
use std::os::unix::process::CommandExt;
use std::path::PathBuf;
use std::process::{Command, Stdio};
use std::sync::Mutex;
use sysinfo::{Pid, Process, Signal, System};
use tracing::{debug, info};
pub struct ProcessController { pub struct ProcessController {
system: System, system: System,
@@ -114,11 +114,16 @@ impl ProcessController {
let intermediate_pid = child.id(); let intermediate_pid = child.id();
// Wait for intermediate process (it will exit immediately after forking) // Wait for intermediate process (it will exit immediately after forking)
child.wait().context("Failed to wait for intermediate process")?; child
.wait()
.context("Failed to wait for intermediate process")?;
// The actual g3 process is now running as orphan // The actual g3 process is now running as orphan
// We need to scan for it by matching workspace and recent start time // We need to scan for it by matching workspace and recent start time
info!("Scanning for newly launched g3 process in workspace: {}", workspace); info!(
"Scanning for newly launched g3 process in workspace: {}",
workspace
);
// Wait even longer for the process to fully start and appear in process list // Wait even longer for the process to fully start and appear in process list
std::thread::sleep(std::time::Duration::from_millis(2500)); std::thread::sleep(std::time::Duration::from_millis(2500));
@@ -151,7 +156,8 @@ impl ProcessController {
if has_workspace { if has_workspace {
// Check if it's recent (started within last 10 seconds) // Check if it's recent (started within last 10 seconds)
let now = std::time::SystemTime::now(); let now = std::time::SystemTime::now();
let start_time = std::time::UNIX_EPOCH + std::time::Duration::from_secs(process.start_time()); let start_time =
std::time::UNIX_EPOCH + std::time::Duration::from_secs(process.start_time());
if let Ok(duration) = now.duration_since(start_time) { if let Ok(duration) = now.duration_since(start_time) {
if duration.as_secs() < 10 { if duration.as_secs() < 10 {
found_pid = Some(pid.as_u32()); found_pid = Some(pid.as_u32());
@@ -283,7 +289,8 @@ impl ProcessController {
} }
// Only return params if we have the minimum required fields // Only return params if we have the minimum required fields
if let (Some(ws), Some(prov), Some(mdl), Some(prmt)) = (workspace, provider, model, prompt) { if let (Some(ws), Some(prov), Some(mdl), Some(prmt)) = (workspace, provider, model, prompt)
{
Some(LaunchParams { Some(LaunchParams {
workspace: ws, workspace: ws,
provider: prov, provider: prov,
+18 -14
View File
@@ -2,7 +2,7 @@ use crate::models::{ExecutionMethod, Instance, InstanceStatus, InstanceType};
use anyhow::Result; use anyhow::Result;
use chrono::{DateTime, Utc}; use chrono::{DateTime, Utc};
use std::path::PathBuf; use std::path::PathBuf;
use sysinfo::{System, Pid, Process}; use sysinfo::{Pid, Process, System};
use tracing::{debug, info, warn}; use tracing::{debug, info, warn};
pub struct ProcessDetector { pub struct ProcessDetector {
@@ -41,12 +41,7 @@ impl ProcessDetector {
Ok(instances) Ok(instances)
} }
fn parse_g3_process( fn parse_g3_process(&self, pid: Pid, process: &Process, cmd: &[String]) -> Option<Instance> {
&self,
pid: Pid,
process: &Process,
cmd: &[String],
) -> Option<Instance> {
let cmd_str = cmd.join(" "); let cmd_str = cmd.join(" ");
// Exclude g3-console itself // Exclude g3-console itself
@@ -55,10 +50,16 @@ impl ProcessDetector {
} }
// Check if this is a g3 binary (more comprehensive check) // Check if this is a g3 binary (more comprehensive check)
let is_g3_binary = cmd.get(0).map(|s| { let is_g3_binary = cmd
(s.ends_with("g3") || s.ends_with("/g3") || s.contains("/target/release/g3") || s.contains("/target/debug/g3")) .get(0)
&& !s.contains("g3-") // Exclude other g3-* binaries .map(|s| {
}).unwrap_or(false); (s.ends_with("g3")
|| s.ends_with("/g3")
|| s.contains("/target/release/g3")
|| s.contains("/target/debug/g3"))
&& !s.contains("g3-") // Exclude other g3-* binaries
})
.unwrap_or(false);
// Check if this is cargo run with g3 (not g3-console or other variants) // Check if this is cargo run with g3 (not g3-console or other variants)
let is_cargo_run = cmd.get(0).map(|s| s.contains("cargo")).unwrap_or(false) let is_cargo_run = cmd.get(0).map(|s| s.contains("cargo")).unwrap_or(false)
@@ -97,8 +98,8 @@ impl ProcessDetector {
let model = self.extract_flag_value(cmd, "--model"); let model = self.extract_flag_value(cmd, "--model");
// Get start time // Get start time
let start_time = DateTime::from_timestamp(process.start_time() as i64, 0) let start_time =
.unwrap_or_else(Utc::now); DateTime::from_timestamp(process.start_time() as i64, 0).unwrap_or_else(Utc::now);
// Generate instance ID from PID and start time // Generate instance ID from PID and start time
let id = format!("{}_{}", pid, start_time.timestamp()); let id = format!("{}_{}", pid, start_time.timestamp());
@@ -158,7 +159,10 @@ impl ProcessDetector {
} }
// Final fallback: use current directory of console // Final fallback: use current directory of console
warn!("Could not determine workspace for PID {}, using current directory", pid); warn!(
"Could not determine workspace for PID {}, using current directory",
pid
);
std::env::current_dir().ok() std::env::current_dir().ok()
} }
+2 -2
View File
@@ -1,5 +1,5 @@
pub mod detector;
pub mod controller; pub mod controller;
pub mod detector;
pub use detector::*;
pub use controller::*; pub use controller::*;
pub use detector::*;
+1 -1
View File
@@ -1,6 +1,6 @@
//! Inspect tree-sitter AST structure for Rust code //! Inspect tree-sitter AST structure for Rust code
use tree_sitter::{Parser, Language}; use tree_sitter::{Language, Parser};
fn print_tree(node: tree_sitter::Node, source: &str, indent: usize) { fn print_tree(node: tree_sitter::Node, source: &str, indent: usize) {
let indent_str = " ".repeat(indent); let indent_str = " ".repeat(indent);
@@ -1,6 +1,6 @@
//! Inspect tree-sitter AST structure for Python code //! Inspect tree-sitter AST structure for Python code
use tree_sitter::{Parser, Language}; use tree_sitter::{Language, Parser};
fn print_tree(node: tree_sitter::Node, source: &str, indent: usize) { fn print_tree(node: tree_sitter::Node, source: &str, indent: usize) {
let indent_str = " ".repeat(indent); let indent_str = " ".repeat(indent);
+1 -1
View File
@@ -1,7 +1,7 @@
//! Test Python async query //! Test Python async query
use tree_sitter::{Parser, Query, QueryCursor, Language};
use streaming_iterator::StreamingIterator; use streaming_iterator::StreamingIterator;
use tree_sitter::{Language, Parser, Query, QueryCursor};
fn main() -> anyhow::Result<()> { fn main() -> anyhow::Result<()> {
let source_code = r#" let source_code = r#"
+9 -10
View File
@@ -3,8 +3,8 @@ use anyhow::{anyhow, Result};
use std::collections::HashMap; use std::collections::HashMap;
use std::fs; use std::fs;
use std::path::Path; use std::path::Path;
use tree_sitter::{Language, Parser, Query, QueryCursor};
use streaming_iterator::StreamingIterator; use streaming_iterator::StreamingIterator;
use tree_sitter::{Language, Parser, Query, QueryCursor};
use walkdir::WalkDir; use walkdir::WalkDir;
pub struct TreeSitterSearcher { pub struct TreeSitterSearcher {
@@ -50,7 +50,8 @@ impl TreeSitterSearcher {
// Create separate parser for "js" alias // Create separate parser for "js" alias
let mut parser_js = Parser::new(); let mut parser_js = Parser::new();
parser_js.set_language(&language) parser_js
.set_language(&language)
.map_err(|e| anyhow!("Failed to set JavaScript language: {}", e))?; .map_err(|e| anyhow!("Failed to set JavaScript language: {}", e))?;
parsers.insert("js".to_string(), parser_js); parsers.insert("js".to_string(), parser_js);
languages.insert("javascript".to_string(), language.clone()); languages.insert("javascript".to_string(), language.clone());
@@ -68,7 +69,8 @@ impl TreeSitterSearcher {
// Create separate parser for "ts" alias // Create separate parser for "ts" alias
let mut parser_ts = Parser::new(); let mut parser_ts = Parser::new();
parser_ts.set_language(&language) parser_ts
.set_language(&language)
.map_err(|e| anyhow!("Failed to set TypeScript language: {}", e))?; .map_err(|e| anyhow!("Failed to set TypeScript language: {}", e))?;
parsers.insert("ts".to_string(), parser_ts); parsers.insert("ts".to_string(), parser_ts);
languages.insert("typescript".to_string(), language.clone()); languages.insert("typescript".to_string(), language.clone());
@@ -215,8 +217,8 @@ impl TreeSitterSearcher {
.ok_or_else(|| anyhow!("Language not found: {}", spec.language))?; .ok_or_else(|| anyhow!("Language not found: {}", spec.language))?;
// Parse query // Parse query
let query = Query::new(language, &spec.query) let query =
.map_err(|e| anyhow!("Invalid query: {}", e))?; Query::new(language, &spec.query).map_err(|e| anyhow!("Invalid query: {}", e))?;
let mut matches = Vec::new(); let mut matches = Vec::new();
let mut files_searched = 0; let mut files_searched = 0;
@@ -255,11 +257,8 @@ impl TreeSitterSearcher {
if let Ok(source_code) = fs::read_to_string(path) { if let Ok(source_code) = fs::read_to_string(path) {
if let Some(tree) = parser.parse(&source_code, None) { if let Some(tree) = parser.parse(&source_code, None) {
let mut cursor = QueryCursor::new(); let mut cursor = QueryCursor::new();
let mut query_matches = cursor.matches( let mut query_matches =
&query, cursor.matches(&query, tree.root_node(), source_code.as_bytes());
tree.root_node(),
source_code.as_bytes(),
);
query_matches.advance(); query_matches.advance();
while let Some(query_match) = query_matches.get() { while let Some(query_match) = query_matches.get() {
+81 -25
View File
@@ -191,23 +191,36 @@ pub fn classify_error(error: &anyhow::Error) -> ErrorType {
let error_str = error.to_string().to_lowercase(); let error_str = error.to_string().to_lowercase();
// Check for recoverable error patterns // Check for recoverable error patterns
if error_str.contains("rate limit") || error_str.contains("rate_limit") || error_str.contains("429") { if error_str.contains("rate limit")
|| error_str.contains("rate_limit")
|| error_str.contains("429")
{
return ErrorType::Recoverable(RecoverableError::RateLimit); return ErrorType::Recoverable(RecoverableError::RateLimit);
} }
if error_str.contains("network") || error_str.contains("connection") || if error_str.contains("network")
error_str.contains("dns") || error_str.contains("refused") { || error_str.contains("connection")
|| error_str.contains("dns")
|| error_str.contains("refused")
{
return ErrorType::Recoverable(RecoverableError::NetworkError); return ErrorType::Recoverable(RecoverableError::NetworkError);
} }
if error_str.contains("500") || error_str.contains("502") || if error_str.contains("500")
error_str.contains("503") || error_str.contains("504") || || error_str.contains("502")
error_str.contains("server error") || error_str.contains("internal error") { || error_str.contains("503")
|| error_str.contains("504")
|| error_str.contains("server error")
|| error_str.contains("internal error")
{
return ErrorType::Recoverable(RecoverableError::ServerError); return ErrorType::Recoverable(RecoverableError::ServerError);
} }
if error_str.contains("busy") || error_str.contains("overloaded") || if error_str.contains("busy")
error_str.contains("capacity") || error_str.contains("unavailable") { || error_str.contains("overloaded")
|| error_str.contains("capacity")
|| error_str.contains("unavailable")
{
return ErrorType::Recoverable(RecoverableError::ModelBusy); return ErrorType::Recoverable(RecoverableError::ModelBusy);
} }
@@ -216,18 +229,24 @@ pub fn classify_error(error: &anyhow::Error) -> ErrorType {
error_str.contains("timed out") || error_str.contains("timed out") ||
error_str.contains("operation timed out") || error_str.contains("operation timed out") ||
error_str.contains("request or response body error") || // Common timeout pattern error_str.contains("request or response body error") || // Common timeout pattern
error_str.contains("stream error") && error_str.contains("timed out") { error_str.contains("stream error") && error_str.contains("timed out")
{
return ErrorType::Recoverable(RecoverableError::Timeout); return ErrorType::Recoverable(RecoverableError::Timeout);
} }
// Check for context length exceeded errors (HTTP 400 with specific messages) // Check for context length exceeded errors (HTTP 400 with specific messages)
if (error_str.contains("400") || error_str.contains("bad request")) && if (error_str.contains("400") || error_str.contains("bad request"))
(error_str.contains("context length") || error_str.contains("prompt is too long") || && (error_str.contains("context length")
error_str.contains("maximum context length") || error_str.contains("context_length_exceeded")) { || error_str.contains("prompt is too long")
|| error_str.contains("maximum context length")
|| error_str.contains("context_length_exceeded"))
{
return ErrorType::Recoverable(RecoverableError::ContextLengthExceeded); return ErrorType::Recoverable(RecoverableError::ContextLengthExceeded);
} }
if error_str.contains("token") && (error_str.contains("limit") || error_str.contains("exceeded")) { if error_str.contains("token")
&& (error_str.contains("limit") || error_str.contains("exceeded"))
{
return ErrorType::Recoverable(RecoverableError::TokenLimit); return ErrorType::Recoverable(RecoverableError::TokenLimit);
} }
@@ -243,7 +262,9 @@ fn calculate_autonomous_retry_delay(attempt: u32) -> Duration {
// Distribute 6 retries over 10 minutes (600 seconds) // Distribute 6 retries over 10 minutes (600 seconds)
// Base delays: 10s, 30s, 60s, 120s, 180s, 200s = 600s total // Base delays: 10s, 30s, 60s, 120s, 180s, 200s = 600s total
let base_delays_ms = [10000, 30000, 60000, 120000, 180000, 200000]; let base_delays_ms = [10000, 30000, 60000, 120000, 180000, 200000];
let base_delay = base_delays_ms.get(attempt.saturating_sub(1) as usize).unwrap_or(&200000); let base_delay = base_delays_ms
.get(attempt.saturating_sub(1) as usize)
.unwrap_or(&200000);
// Add jitter of ±30% to prevent thundering herd // Add jitter of ±30% to prevent thundering herd
let jitter = (*base_delay as f64 * 0.3 * rng.gen::<f64>()) as u64; let jitter = (*base_delay as f64 * 0.3 * rng.gen::<f64>()) as u64;
@@ -263,7 +284,11 @@ pub fn calculate_retry_delay(attempt: u32, is_autonomous: bool) -> Duration {
} }
use rand::Rng; use rand::Rng;
let max_retry_delay_ms = if is_autonomous { AUTONOMOUS_MAX_RETRY_DELAY_MS } else { DEFAULT_MAX_RETRY_DELAY_MS }; let max_retry_delay_ms = if is_autonomous {
AUTONOMOUS_MAX_RETRY_DELAY_MS
} else {
DEFAULT_MAX_RETRY_DELAY_MS
};
// Exponential backoff: delay = base * 2^attempt // Exponential backoff: delay = base * 2^attempt
let base_delay = BASE_RETRY_DELAY_MS * (2_u64.pow(attempt.saturating_sub(1))); let base_delay = BASE_RETRY_DELAY_MS * (2_u64.pow(attempt.saturating_sub(1)));
@@ -370,7 +395,11 @@ fn truncate_for_logging(s: &str, max_len: usize) -> String {
truncate_at = max_len.min(s.len()); truncate_at = max_len.min(s.len());
} }
format!("{}... (truncated, {} total bytes)", &s[..truncate_at], s.len()) format!(
"{}... (truncated, {} total bytes)",
&s[..truncate_at],
s.len()
)
} }
} }
@@ -398,37 +427,64 @@ mod tests {
fn test_error_classification() { fn test_error_classification() {
// Rate limit errors // Rate limit errors
let error = anyhow!("Rate limit exceeded"); let error = anyhow!("Rate limit exceeded");
assert_eq!(classify_error(&error), ErrorType::Recoverable(RecoverableError::RateLimit)); assert_eq!(
classify_error(&error),
ErrorType::Recoverable(RecoverableError::RateLimit)
);
let error = anyhow!("HTTP 429 Too Many Requests"); let error = anyhow!("HTTP 429 Too Many Requests");
assert_eq!(classify_error(&error), ErrorType::Recoverable(RecoverableError::RateLimit)); assert_eq!(
classify_error(&error),
ErrorType::Recoverable(RecoverableError::RateLimit)
);
// Network errors // Network errors
let error = anyhow!("Network connection failed"); let error = anyhow!("Network connection failed");
assert_eq!(classify_error(&error), ErrorType::Recoverable(RecoverableError::NetworkError)); assert_eq!(
classify_error(&error),
ErrorType::Recoverable(RecoverableError::NetworkError)
);
// Server errors // Server errors
let error = anyhow!("HTTP 503 Service Unavailable"); let error = anyhow!("HTTP 503 Service Unavailable");
assert_eq!(classify_error(&error), ErrorType::Recoverable(RecoverableError::ServerError)); assert_eq!(
classify_error(&error),
ErrorType::Recoverable(RecoverableError::ServerError)
);
// Model busy // Model busy
let error = anyhow!("Model is busy, please try again"); let error = anyhow!("Model is busy, please try again");
assert_eq!(classify_error(&error), ErrorType::Recoverable(RecoverableError::ModelBusy)); assert_eq!(
classify_error(&error),
ErrorType::Recoverable(RecoverableError::ModelBusy)
);
// Timeout // Timeout
let error = anyhow!("Request timed out"); let error = anyhow!("Request timed out");
assert_eq!(classify_error(&error), ErrorType::Recoverable(RecoverableError::Timeout)); assert_eq!(
classify_error(&error),
ErrorType::Recoverable(RecoverableError::Timeout)
);
// Token limit // Token limit
let error = anyhow!("Token limit exceeded"); let error = anyhow!("Token limit exceeded");
assert_eq!(classify_error(&error), ErrorType::Recoverable(RecoverableError::TokenLimit)); assert_eq!(
classify_error(&error),
ErrorType::Recoverable(RecoverableError::TokenLimit)
);
// Context length exceeded // Context length exceeded
let error = anyhow!("HTTP 400 Bad Request: context length exceeded"); let error = anyhow!("HTTP 400 Bad Request: context length exceeded");
assert_eq!(classify_error(&error), ErrorType::Recoverable(RecoverableError::ContextLengthExceeded)); assert_eq!(
classify_error(&error),
ErrorType::Recoverable(RecoverableError::ContextLengthExceeded)
);
let error = anyhow!("Error 400: prompt is too long"); let error = anyhow!("Error 400: prompt is too long");
assert_eq!(classify_error(&error), ErrorType::Recoverable(RecoverableError::ContextLengthExceeded)); assert_eq!(
classify_error(&error),
ErrorType::Recoverable(RecoverableError::ContextLengthExceeded)
);
// Non-recoverable // Non-recoverable
let error = anyhow!("Invalid API key"); let error = anyhow!("Invalid API key");
+4 -4
View File
@@ -17,7 +17,7 @@ mod tests {
"test prompt".to_string(), "test prompt".to_string(),
None, None,
100, 100,
false, // quiet parameter false, // quiet parameter
); );
let result = retry_with_backoff( let result = retry_with_backoff(
@@ -57,7 +57,7 @@ mod tests {
"test prompt".to_string(), "test prompt".to_string(),
None, None,
100, 100,
false, // quiet parameter false, // quiet parameter
); );
let result: Result<&str, _> = retry_with_backoff( let result: Result<&str, _> = retry_with_backoff(
@@ -91,7 +91,7 @@ mod tests {
"test prompt".to_string(), "test prompt".to_string(),
None, None,
100, 100,
false, // quiet parameter false, // quiet parameter
); );
let result: Result<&str, _> = retry_with_backoff( let result: Result<&str, _> = retry_with_backoff(
@@ -124,7 +124,7 @@ mod tests {
long_prompt, long_prompt,
None, None,
100, 100,
false, // quiet parameter false, // quiet parameter
); );
// The prompt should be truncated to 1000 chars // The prompt should be truncated to 1000 chars
+1 -1
View File
@@ -29,7 +29,7 @@ struct FixedJsonToolState {
brace_depth: i32, brace_depth: i32,
buffer: String, buffer: String,
json_start_in_buffer: Option<usize>, // Position where confirmed JSON tool call starts json_start_in_buffer: Option<usize>, // Position where confirmed JSON tool call starts
content_returned_up_to: usize, // Track how much content we've already returned content_returned_up_to: usize, // Track how much content we've already returned
potential_json_start: Option<usize>, // Where the potential JSON started potential_json_start: Option<usize>, // Where the potential JSON started
} }
+2 -2
View File
@@ -358,8 +358,8 @@ More text"#;
// 2. Then the same complete JSON appears // 2. Then the same complete JSON appears
let chunks = vec![ let chunks = vec![
"Some text\n", "Some text\n",
r#"{"tool": "str_replace", "args": {"diff":"...","file_path":"./crates/g3-cli"#, // Truncated r#"{"tool": "str_replace", "args": {"diff":"...","file_path":"./crates/g3-cli"#, // Truncated
r#"{"tool": "str_replace", "args": {"diff":"...","file_path":"./crates/g3-cli/src/lib.rs"}}"#, // Complete r#"{"tool": "str_replace", "args": {"diff":"...","file_path":"./crates/g3-cli/src/lib.rs"}}"#, // Complete
"\nMore text", "\nMore text",
]; ];
+468 -196
View File
File diff suppressed because it is too large Load Diff
+4 -1
View File
@@ -58,7 +58,10 @@ impl Project {
} }
/// Create a project for autonomous mode with requirements text override /// Create a project for autonomous mode with requirements text override
pub fn new_autonomous_with_requirements(workspace_dir: PathBuf, requirements_text: String) -> Result<Self> { pub fn new_autonomous_with_requirements(
workspace_dir: PathBuf,
requirements_text: String,
) -> Result<Self> {
let mut project = Self::new(workspace_dir.clone()); let mut project = Self::new(workspace_dir.clone());
project.autonomous = true; project.autonomous = true;
project.requirements_text = Some(requirements_text); project.requirements_text = Some(requirements_text);
+1 -1
View File
@@ -189,7 +189,7 @@ Do not explain what you're going to do - just do it by calling the tools.
"; ";
pub const SYSTEM_PROMPT_FOR_NATIVE_TOOL_USE: &'static str = pub const SYSTEM_PROMPT_FOR_NATIVE_TOOL_USE: &'static str =
concatcp!(SYSTEM_NATIVE_TOOL_CALLS, CODING_STYLE); concatcp!(SYSTEM_NATIVE_TOOL_CALLS, CODING_STYLE);
/// Generate system prompt based on whether multiple tool calls are allowed /// Generate system prompt based on whether multiple tool calls are allowed
pub fn get_system_prompt_for_native(allow_multiple: bool) -> String { pub fn get_system_prompt_for_native(allow_multiple: bool) -> String {
+12 -5
View File
@@ -67,7 +67,8 @@ impl TaskResult {
let blocks: Vec<&str> = content_without_timing.split("\n\n").collect(); let blocks: Vec<&str> = content_without_timing.split("\n\n").collect();
// Find the last non-empty block that isn't just whitespace // Find the last non-empty block that isn't just whitespace
blocks.iter() blocks
.iter()
.rev() .rev()
.find(|block| !block.trim().is_empty()) .find(|block| !block.trim().is_empty())
.map(|block| block.trim().to_string()) .map(|block| block.trim().to_string())
@@ -79,7 +80,8 @@ impl TaskResult {
/// Check if the response contains an approval (for autonomous mode) /// Check if the response contains an approval (for autonomous mode)
pub fn is_approved(&self) -> bool { pub fn is_approved(&self) -> bool {
self.extract_final_output().contains("IMPLEMENTATION_APPROVED") self.extract_final_output()
.contains("IMPLEMENTATION_APPROVED")
} }
} }
@@ -91,7 +93,8 @@ mod tests {
fn test_extract_last_block() { fn test_extract_last_block() {
// Test case 1: Response with timing info // Test case 1: Response with timing info
let context_window = ContextWindow::new(1000); let context_window = ContextWindow::new(1000);
let response_with_timing = "Some initial content\n\nFinal block content\n\n⏱️ 2.3s | 💭 1.2s".to_string(); let response_with_timing =
"Some initial content\n\nFinal block content\n\n⏱️ 2.3s | 💭 1.2s".to_string();
let result = TaskResult::new(response_with_timing, context_window.clone()); let result = TaskResult::new(response_with_timing, context_window.clone());
assert_eq!(result.extract_last_block(), "Final block content"); assert_eq!(result.extract_last_block(), "Final block content");
@@ -138,10 +141,14 @@ mod tests {
// Test case 1: Response with final_output tool call // Test case 1: Response with final_output tool call
let response_with_final_output = "Analyzing files...\n\nCalling final_output\n\nThis is the complete feedback\nwith multiple lines\nand important details\n\n⏱️ 2.3s".to_string(); let response_with_final_output = "Analyzing files...\n\nCalling final_output\n\nThis is the complete feedback\nwith multiple lines\nand important details\n\n⏱️ 2.3s".to_string();
let result = TaskResult::new(response_with_final_output, context_window.clone()); let result = TaskResult::new(response_with_final_output, context_window.clone());
assert_eq!(result.extract_final_output(), "This is the complete feedback\nwith multiple lines\nand important details"); assert_eq!(
result.extract_final_output(),
"This is the complete feedback\nwith multiple lines\nand important details"
);
// Test case 2: Response with IMPLEMENTATION_APPROVED in final_output // Test case 2: Response with IMPLEMENTATION_APPROVED in final_output
let response_approved = "Review complete\n\nfinal_output called\n\nIMPLEMENTATION_APPROVED".to_string(); let response_approved =
"Review complete\n\nfinal_output called\n\nIMPLEMENTATION_APPROVED".to_string();
let result = TaskResult::new(response_approved, context_window.clone()); let result = TaskResult::new(response_approved, context_window.clone());
assert_eq!(result.extract_final_output(), "IMPLEMENTATION_APPROVED"); assert_eq!(result.extract_final_output(), "IMPLEMENTATION_APPROVED");
assert!(result.is_approved()); assert!(result.is_approved());
@@ -6,10 +6,14 @@ use std::sync::Arc;
fn test_task_result_basic_functionality() { fn test_task_result_basic_functionality() {
// Create a context window with some messages // Create a context window with some messages
let mut context = ContextWindow::new(10000); let mut context = ContextWindow::new(10000);
context.add_message(Message::new(MessageRole::User, "Test message 1".to_string()) context.add_message(Message::new(
); MessageRole::User,
context.add_message(Message::new(MessageRole::Assistant, "Response 1".to_string()) "Test message 1".to_string(),
); ));
context.add_message(Message::new(
MessageRole::Assistant,
"Response 1".to_string(),
));
// Create a TaskResult // Create a TaskResult
let response = "This is the response\n\nFinal output block".to_string(); let response = "This is the response\n\nFinal output block".to_string();
@@ -70,7 +74,11 @@ fn test_is_approved_detection() {
for response in approved_responses { for response in approved_responses {
let result = TaskResult::new(response.to_string(), context.clone()); let result = TaskResult::new(response.to_string(), context.clone());
assert!(result.is_approved(), "Failed to detect approval in: {}", response); assert!(
result.is_approved(),
"Failed to detect approval in: {}",
response
);
} }
// Test not approved cases // Test not approved cases
@@ -84,7 +92,11 @@ fn test_is_approved_detection() {
for response in not_approved_responses { for response in not_approved_responses {
let result = TaskResult::new(response.to_string(), context.clone()); let result = TaskResult::new(response.to_string(), context.clone());
assert!(!result.is_approved(), "Incorrectly detected approval in: {}", response); assert!(
!result.is_approved(),
"Incorrectly detected approval in: {}",
response
);
} }
} }
@@ -96,7 +108,14 @@ fn test_context_window_preservation() {
// Add some messages // Add some messages
for i in 0..5 { for i in 0..5 {
context.add_message(Message::new(if i % 2 == 0 { MessageRole::User } else { MessageRole::Assistant }, format!("Message {}", i))); context.add_message(Message::new(
if i % 2 == 0 {
MessageRole::User
} else {
MessageRole::Assistant
},
format!("Message {}", i),
));
} }
// Create TaskResult // Create TaskResult
@@ -109,10 +128,16 @@ fn test_context_window_preservation() {
// Verify messages are preserved correctly // Verify messages are preserved correctly
for i in 0..5 { for i in 0..5 {
let is_user = matches!(result.context_window.conversation_history[i].role, MessageRole::User); let is_user = matches!(
result.context_window.conversation_history[i].role,
MessageRole::User
);
let expected_is_user = i % 2 == 0; let expected_is_user = i % 2 == 0;
assert_eq!(is_user, expected_is_user, "Message {} has wrong role", i); assert_eq!(is_user, expected_is_user, "Message {} has wrong role", i);
assert_eq!(result.context_window.conversation_history[i].content, format!("Message {}", i)); assert_eq!(
result.context_window.conversation_history[i].content,
format!("Message {}", i)
);
} }
} }
@@ -181,7 +206,8 @@ fn test_edge_cases_and_special_characters() {
assert_eq!(result.extract_last_block(), "✅ Final part with emojis 🎉"); assert_eq!(result.extract_last_block(), "✅ Final part with emojis 🎉");
// Test with code blocks // Test with code blocks
let response_with_code = "Explanation\n\n```rust\nfn main() {}\n```\n\nFinal comment".to_string(); let response_with_code =
"Explanation\n\n```rust\nfn main() {}\n```\n\nFinal comment".to_string();
let result = TaskResult::new(response_with_code, context.clone()); let result = TaskResult::new(response_with_code, context.clone());
assert_eq!(result.extract_last_block(), "Final comment"); assert_eq!(result.extract_last_block(), "Final comment");
@@ -203,7 +229,10 @@ fn test_large_response_handling() {
large_response.push_str("This is the final block after 100 other blocks"); large_response.push_str("This is the final block after 100 other blocks");
let result = TaskResult::new(large_response, context); let result = TaskResult::new(large_response, context);
assert_eq!(result.extract_last_block(), "This is the final block after 100 other blocks"); assert_eq!(
result.extract_last_block(),
"This is the final block after 100 other blocks"
);
} }
#[test] #[test]
@@ -237,4 +266,3 @@ fn test_concurrent_access() {
handle.join().unwrap(); handle.join().unwrap();
} }
} }
+12 -4
View File
@@ -55,7 +55,9 @@ pub trait UiWriter: Send + Sync {
/// Returns true if this UI writer wants full, untruncated output /// Returns true if this UI writer wants full, untruncated output
/// Default is false (truncate for human readability) /// Default is false (truncate for human readability)
fn wants_full_output(&self) -> bool { false } fn wants_full_output(&self) -> bool {
false
}
/// Prompt the user for a yes/no confirmation /// Prompt the user for a yes/no confirmation
fn prompt_user_yes_no(&self, message: &str) -> bool; fn prompt_user_yes_no(&self, message: &str) -> bool;
@@ -86,7 +88,13 @@ impl UiWriter for NullUiWriter {
fn print_agent_response(&self, _content: &str) {} fn print_agent_response(&self, _content: &str) {}
fn notify_sse_received(&self) {} fn notify_sse_received(&self) {}
fn flush(&self) {} fn flush(&self) {}
fn wants_full_output(&self) -> bool { false } fn wants_full_output(&self) -> bool {
fn prompt_user_yes_no(&self, _message: &str) -> bool { true } false
fn prompt_user_choice(&self, _message: &str, _options: &[&str]) -> usize { 0 } }
fn prompt_user_yes_no(&self, _message: &str) -> bool {
true
}
fn prompt_user_choice(&self, _message: &str, _options: &[&str]) -> usize {
0
}
} }
+43 -12
View File
@@ -47,7 +47,10 @@ pub async fn another_async(x: i32) -> Result<(), ()> {
assert_eq!(response.searches.len(), 1); assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0]; let search_result = &response.searches[0];
assert_eq!(search_result.name, "find_async_functions"); assert_eq!(search_result.name, "find_async_functions");
assert_eq!(search_result.match_count, 2, "Should find 2 async functions"); assert_eq!(
search_result.match_count, 2,
"Should find 2 async functions"
);
assert!(search_result.error.is_none()); assert!(search_result.error.is_none());
// Check that we found the right functions // Check that we found the right functions
@@ -107,7 +110,10 @@ pub async fn another_async(x: i32) -> Result<(), ()> {
assert_eq!(response.searches.len(), 1); assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0]; let search_result = &response.searches[0];
assert_eq!(search_result.name, "find_all_functions"); assert_eq!(search_result.name, "find_all_functions");
assert_eq!(search_result.match_count, 3, "Should find 3 functions total"); assert_eq!(
search_result.match_count, 3,
"Should find 3 functions total"
);
assert!(search_result.error.is_none()); assert!(search_result.error.is_none());
// Check that we found all functions // Check that we found all functions
@@ -229,10 +235,16 @@ pub fn target_function() {
let context = match_result.context.as_ref().unwrap(); let context = match_result.context.as_ref().unwrap();
assert!(context.contains("Line 2"), "Should include 2 lines before"); assert!(context.contains("Line 2"), "Should include 2 lines before");
assert!(context.contains("target_function"), "Should include the function"); assert!(
context.contains("target_function"),
"Should include the function"
);
// Note: context_lines=2 means 2 lines before and after the match line (line 4) // Note: context_lines=2 means 2 lines before and after the match line (line 4)
// So we get lines 2-6, which includes up to println but not the closing brace // So we get lines 2-6, which includes up to println but not the closing brace
assert!(context.contains("println"), "Should include 2 lines after the match"); assert!(
context.contains("println"),
"Should include 2 lines after the match"
);
// Cleanup // Cleanup
fs::remove_dir_all(&test_dir).ok(); fs::remove_dir_all(&test_dir).ok();
@@ -338,7 +350,10 @@ class MyClass:
assert_eq!(response.searches.len(), 1); assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0]; let search_result = &response.searches[0];
assert_eq!(search_result.match_count, 3, "Should find 3 functions in Python (2 regular + 1 async + 1 method)"); assert_eq!(
search_result.match_count, 3,
"Should find 3 functions in Python (2 regular + 1 async + 1 method)"
);
let function_names: Vec<String> = search_result let function_names: Vec<String> = search_result
.matches .matches
@@ -396,7 +411,10 @@ class MyClass {
assert_eq!(response.searches.len(), 1); assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0]; let search_result = &response.searches[0];
assert_eq!(search_result.match_count, 2, "Should find 2 functions in JavaScript"); assert_eq!(
search_result.match_count, 2,
"Should find 2 functions in JavaScript"
);
let function_names: Vec<String> = search_result let function_names: Vec<String> = search_result
.matches .matches
@@ -439,10 +457,15 @@ async fn test_go_search() {
eprintln!("Go search result: {:?}", response.searches[0]); eprintln!("Go search result: {:?}", response.searches[0]);
eprintln!("Match count: {}", response.searches[0].matches.len()); eprintln!("Match count: {}", response.searches[0].matches.len());
eprintln!("Error: {:?}", response.searches[0].error); eprintln!("Error: {:?}", response.searches[0].error);
assert!(response.searches[0].matches.len() > 0, "No matches found for Go search"); assert!(
response.searches[0].matches.len() > 0,
"No matches found for Go search"
);
// Should find main and greet functions // Should find main and greet functions
let names: Vec<&str> = response.searches[0].matches.iter() let names: Vec<&str> = response.searches[0]
.matches
.iter()
.filter_map(|m| m.captures.get("name").map(|s| s.as_str())) .filter_map(|m| m.captures.get("name").map(|s| s.as_str()))
.collect(); .collect();
assert!(names.contains(&"main")); assert!(names.contains(&"main"));
@@ -476,7 +499,9 @@ async fn test_java_search() {
assert!(response.searches[0].matches.len() > 0); assert!(response.searches[0].matches.len() > 0);
// Should find Example class // Should find Example class
let names: Vec<&str> = response.searches[0].matches.iter() let names: Vec<&str> = response.searches[0]
.matches
.iter()
.filter_map(|m| m.captures.get("name").map(|s| s.as_str())) .filter_map(|m| m.captures.get("name").map(|s| s.as_str()))
.collect(); .collect();
assert!(names.contains(&"Example")); assert!(names.contains(&"Example"));
@@ -509,7 +534,9 @@ async fn test_c_search() {
assert!(response.searches[0].matches.len() > 0); assert!(response.searches[0].matches.len() > 0);
// Should find greet, add, and main functions // Should find greet, add, and main functions
let names: Vec<&str> = response.searches[0].matches.iter() let names: Vec<&str> = response.searches[0]
.matches
.iter()
.filter_map(|m| m.captures.get("name").map(|s| s.as_str())) .filter_map(|m| m.captures.get("name").map(|s| s.as_str()))
.collect(); .collect();
assert!(names.contains(&"greet")); assert!(names.contains(&"greet"));
@@ -544,7 +571,9 @@ async fn test_cpp_search() {
assert!(response.searches[0].matches.len() > 0); assert!(response.searches[0].matches.len() > 0);
// Should find Person class // Should find Person class
let names: Vec<&str> = response.searches[0].matches.iter() let names: Vec<&str> = response.searches[0]
.matches
.iter()
.filter_map(|m| m.captures.get("name").map(|s| s.as_str())) .filter_map(|m| m.captures.get("name").map(|s| s.as_str()))
.collect(); .collect();
assert!(names.contains(&"Person")); assert!(names.contains(&"Person"));
@@ -570,7 +599,9 @@ async fn test_kotlin_search() {
assert!(response.searches[0].matches.len() > 0); assert!(response.searches[0].matches.len() > 0);
// Should find Person class // Should find Person class
let names: Vec<&str> = response.searches[0].matches.iter() let names: Vec<&str> = response.searches[0]
.matches
.iter()
.filter_map(|m| m.captures.get("name").map(|s| s.as_str())) .filter_map(|m| m.captures.get("name").map(|s| s.as_str()))
.collect(); .collect();
assert!(names.contains(&"Person")); assert!(names.contains(&"Person"));
+16 -8
View File
@@ -63,10 +63,7 @@ fn test_thin_context_basic() {
format!("Tool result: small result {}", i) format!("Tool result: small result {}", i)
}; };
context.add_message(Message::new( context.add_message(Message::new(MessageRole::User, content));
MessageRole::User,
content,
));
} }
} }
@@ -77,7 +74,11 @@ fn test_thin_context_basic() {
println!("Thinning summary: {}", summary); println!("Thinning summary: {}", summary);
// Should have thinned at least 1 large tool result in the first third // Should have thinned at least 1 large tool result in the first third
assert!(summary.contains("1 tool result"), "Summary was: {}", summary); assert!(
summary.contains("1 tool result"),
"Summary was: {}",
summary
);
assert!(summary.contains("50%")); assert!(summary.contains("50%"));
// Check that the large tool results were replaced // Check that the large tool results were replaced
@@ -160,7 +161,11 @@ fn test_thin_str_replace_tool_calls() {
)); ));
// Add an assistant message with a str_replace tool call containing large diff // Add an assistant message with a str_replace tool call containing large diff
let large_diff = format!("--- old\n{}\n+++ new\n{}", "-old line\n".repeat(100), "+new line\n".repeat(100)); let large_diff = format!(
"--- old\n{}\n+++ new\n{}",
"-old line\n".repeat(100),
"+new line\n".repeat(100)
);
let tool_call_json = format!( let tool_call_json = format!(
r#"{{"tool": "str_replace", "args": {{"file_path": "test.txt", "diff": "{}"}}}}"#, r#"{{"tool": "str_replace", "args": {{"file_path": "test.txt", "diff": "{}"}}}}"#,
large_diff.replace('\n', "\\n") large_diff.replace('\n', "\\n")
@@ -261,8 +266,11 @@ fn test_thin_context_only_affects_first_third() {
if matches!(msg.role, MessageRole::User) && msg.content.starts_with("Tool result:") { if matches!(msg.role, MessageRole::User) && msg.content.starts_with("Tool result:") {
// These should still be large (not thinned) // These should still be large (not thinned)
if i % 2 == 1 { if i % 2 == 1 {
assert!(msg.content.len() > 1000, assert!(
"Message at index {} should not have been thinned", i); msg.content.len() > 1000,
"Message at index {} should not have been thinned",
i
);
} }
} }
} }
@@ -8,7 +8,10 @@ fn test_todo_read_results_not_thinned() {
let mut context = ContextWindow::new(10000); let mut context = ContextWindow::new(10000);
// Add a todo_read tool call // Add a todo_read tool call
context.add_message(Message::new(MessageRole::Assistant, r#"{"tool": "todo_read", "args": {}}"#.to_string())); context.add_message(Message::new(
MessageRole::Assistant,
r#"{"tool": "todo_read", "args": {}}"#.to_string(),
));
// Add a large TODO result (> 500 chars) // Add a large TODO result (> 500 chars)
let large_todo_result = format!( let large_todo_result = format!(
@@ -19,7 +22,10 @@ fn test_todo_read_results_not_thinned() {
// Add more messages to ensure we have enough for "first third" logic // Add more messages to ensure we have enough for "first third" logic
for i in 0..6 { for i in 0..6 {
context.add_message(Message::new(MessageRole::Assistant, format!("Response {}", i))) context.add_message(Message::new(
MessageRole::Assistant,
format!("Response {}", i),
))
} }
// Trigger thinning at 50% // Trigger thinning at 50%
@@ -56,7 +62,13 @@ fn test_todo_write_results_not_thinned() {
// Add a todo_write tool call // Add a todo_write tool call
let large_content = "- [ ] Task\n".repeat(100); let large_content = "- [ ] Task\n".repeat(100);
context.add_message(Message::new(MessageRole::Assistant, format!(r#"{{"tool": "todo_write", "args": {{"content": "{}"}}}}"#, large_content))); context.add_message(Message::new(
MessageRole::Assistant,
format!(
r#"{{"tool": "todo_write", "args": {{"content": "{}"}}}}"#,
large_content
),
));
// Add a large TODO write result // Add a large TODO write result
let large_todo_result = format!( let large_todo_result = format!(
@@ -67,7 +79,10 @@ fn test_todo_write_results_not_thinned() {
// Add more messages // Add more messages
for i in 0..6 { for i in 0..6 {
context.add_message(Message::new(MessageRole::Assistant, format!("Response {}", i))) context.add_message(Message::new(
MessageRole::Assistant,
format!("Response {}", i),
))
} }
// Trigger thinning at 50% // Trigger thinning at 50%
@@ -101,7 +116,10 @@ fn test_non_todo_results_still_thinned() {
let mut context = ContextWindow::new(10000); let mut context = ContextWindow::new(10000);
// Add a non-TODO tool call (e.g., read_file) // Add a non-TODO tool call (e.g., read_file)
context.add_message(Message::new(MessageRole::Assistant, r#"{"tool": "read_file", "args": {"file_path": "test.txt"}}"#.to_string())); context.add_message(Message::new(
MessageRole::Assistant,
r#"{"tool": "read_file", "args": {"file_path": "test.txt"}}"#.to_string(),
));
// Add a large read_file result (> 500 chars) // Add a large read_file result (> 500 chars)
let large_result = format!("Tool result: {}", "x".repeat(1500)); let large_result = format!("Tool result: {}", "x".repeat(1500));
@@ -109,7 +127,10 @@ fn test_non_todo_results_still_thinned() {
// Add more messages // Add more messages
for i in 0..6 { for i in 0..6 {
context.add_message(Message::new(MessageRole::Assistant, format!("Response {}", i))) context.add_message(Message::new(
MessageRole::Assistant,
format!("Response {}", i),
))
} }
// Trigger thinning at 50% // Trigger thinning at 50%
@@ -145,18 +166,21 @@ fn test_todo_read_with_spaces_in_tool_name() {
let mut context = ContextWindow::new(10000); let mut context = ContextWindow::new(10000);
// Add a todo_read tool call with spaces (JSON formatting variation) // Add a todo_read tool call with spaces (JSON formatting variation)
context.add_message(Message::new(MessageRole::Assistant, r#"{"tool": "todo_read", "args": {}}"#.to_string())); context.add_message(Message::new(
MessageRole::Assistant,
r#"{"tool": "todo_read", "args": {}}"#.to_string(),
));
// Add a large TODO result // Add a large TODO result
let large_todo_result = format!( let large_todo_result = format!("Tool result: 📝 TODO list:\n{}", "- [ ] Task\n".repeat(50));
"Tool result: 📝 TODO list:\n{}",
"- [ ] Task\n".repeat(50)
);
context.add_message(Message::new(MessageRole::User, large_todo_result.clone())); context.add_message(Message::new(MessageRole::User, large_todo_result.clone()));
// Add more messages // Add more messages
for i in 0..6 { for i in 0..6 {
context.add_message(Message::new(MessageRole::Assistant, format!("Response {}", i))) context.add_message(Message::new(
MessageRole::Assistant,
format!("Response {}", i),
))
} }
// Trigger thinning // Trigger thinning
+74 -17
View File
@@ -1,11 +1,10 @@
use g3_core::Agent;
use g3_core::ui_writer::NullUiWriter; use g3_core::ui_writer::NullUiWriter;
use g3_core::Agent;
use serial_test::serial; use serial_test::serial;
use std::fs; use std::fs;
use std::path::PathBuf; use std::path::PathBuf;
use tempfile::TempDir; use tempfile::TempDir;
/// Helper to create a test agent in a temporary directory /// Helper to create a test agent in a temporary directory
async fn create_test_agent_in_dir(temp_dir: &TempDir) -> Agent<NullUiWriter> { async fn create_test_agent_in_dir(temp_dir: &TempDir) -> Agent<NullUiWriter> {
// Change to temp directory // Change to temp directory
@@ -46,7 +45,11 @@ async fn test_todo_write_creates_file() {
// Should report success // Should report success
assert!(result.contains("✅"), "Should report success: {}", result); assert!(result.contains("✅"), "Should report success: {}", result);
assert!(result.contains("todo.g3.md"), "Should mention todo.g3.md: {}", result); assert!(
result.contains("todo.g3.md"),
"Should mention todo.g3.md: {}",
result
);
// File should now exist // File should now exist
assert!(todo_path.exists(), "todo.g3.md should exist after write"); assert!(todo_path.exists(), "todo.g3.md should exist after write");
@@ -79,9 +82,21 @@ async fn test_todo_read_from_file() {
let result = agent.execute_tool(&tool_call).await.unwrap(); let result = agent.execute_tool(&tool_call).await.unwrap();
// Should contain the TODO content // Should contain the TODO content
assert!(result.contains("📝 TODO list:"), "Should have TODO list header: {}", result); assert!(
assert!(result.contains("First task"), "Should contain first task: {}", result); result.contains("📝 TODO list:"),
assert!(result.contains("Completed task"), "Should contain completed task: {}", result); "Should have TODO list header: {}",
result
);
assert!(
result.contains("First task"),
"Should contain first task: {}",
result
);
assert!(
result.contains("Completed task"),
"Should contain completed task: {}",
result
);
} }
#[tokio::test] #[tokio::test]
@@ -122,7 +137,10 @@ async fn test_todo_persistence_across_agents() {
} }
// Verify file exists // Verify file exists
assert!(todo_path.exists(), "todo.g3.md should persist after agent drops"); assert!(
todo_path.exists(),
"todo.g3.md should persist after agent drops"
);
// Agent 2: Read TODO (new agent instance) // Agent 2: Read TODO (new agent instance)
{ {
@@ -134,8 +152,16 @@ async fn test_todo_persistence_across_agents() {
let result = agent.execute_tool(&tool_call).await.unwrap(); let result = agent.execute_tool(&tool_call).await.unwrap();
// Should read the persisted content // Should read the persisted content
assert!(result.contains("Persistent task"), "Should read persisted task: {}", result); assert!(
assert!(result.contains("Done task"), "Should read done task: {}", result); result.contains("Persistent task"),
"Should read persisted task: {}",
result
);
assert!(
result.contains("Done task"),
"Should read done task: {}",
result
);
} }
} }
@@ -179,7 +205,10 @@ async fn test_todo_handles_large_content() {
// Create a large TODO (but under the 50k limit) // Create a large TODO (but under the 50k limit)
let mut large_content = String::from("# Large TODO\n\n"); let mut large_content = String::from("# Large TODO\n\n");
for i in 0..100 { for i in 0..100 {
large_content.push_str(&format!("- [ ] Task {} with a long description that exceeds normal line lengths\n", i)); large_content.push_str(&format!(
"- [ ] Task {} with a long description that exceeds normal line lengths\n",
i
));
} }
let tool_call = g3_core::ToolCall { let tool_call = g3_core::ToolCall {
@@ -190,7 +219,11 @@ async fn test_todo_handles_large_content() {
}; };
let result = agent.execute_tool(&tool_call).await.unwrap(); let result = agent.execute_tool(&tool_call).await.unwrap();
assert!(result.contains("✅"), "Should handle large content: {}", result); assert!(
result.contains("✅"),
"Should handle large content: {}",
result
);
// Verify file contains all content // Verify file contains all content
let file_content = fs::read_to_string(&todo_path).unwrap(); let file_content = fs::read_to_string(&todo_path).unwrap();
@@ -217,8 +250,16 @@ async fn test_todo_respects_size_limit() {
let result = agent.execute_tool(&tool_call).await.unwrap(); let result = agent.execute_tool(&tool_call).await.unwrap();
// Should reject content that's too large // Should reject content that's too large
assert!(result.contains("❌"), "Should reject oversized content: {}", result); assert!(
assert!(result.contains("too large"), "Should mention size limit: {}", result); result.contains("❌"),
"Should reject oversized content: {}",
result
);
assert!(
result.contains("too large"),
"Should mention size limit: {}",
result
);
} }
#[tokio::test] #[tokio::test]
@@ -241,7 +282,11 @@ async fn test_todo_agent_initialization_loads_file() {
}; };
let result = agent.execute_tool(&tool_call).await.unwrap(); let result = agent.execute_tool(&tool_call).await.unwrap();
assert!(result.contains("Pre-existing task"), "Should load file on init: {}", result); assert!(
result.contains("Pre-existing task"),
"Should load file on init: {}",
result
);
} }
#[tokio::test] #[tokio::test]
@@ -274,9 +319,17 @@ async fn test_todo_handles_unicode_content() {
}; };
let result = agent.execute_tool(&read_call).await.unwrap(); let result = agent.execute_tool(&read_call).await.unwrap();
assert!(result.contains("日本語"), "Should preserve Japanese: {}", result); assert!(
result.contains("日本語"),
"Should preserve Japanese: {}",
result
);
assert!(result.contains("🚀"), "Should preserve emoji: {}", result); assert!(result.contains("🚀"), "Should preserve emoji: {}", result);
assert!(result.contains("Ελληνικά"), "Should preserve Greek: {}", result); assert!(
result.contains("Ελληνικά"),
"Should preserve Greek: {}",
result
);
} }
#[tokio::test] #[tokio::test]
@@ -327,5 +380,9 @@ async fn test_todo_whitespace_only_content() {
let result = agent.execute_tool(&read_call).await.unwrap(); let result = agent.execute_tool(&read_call).await.unwrap();
// Should report as empty (whitespace is trimmed) // Should report as empty (whitespace is trimmed)
assert!(result.contains("empty"), "Whitespace-only should be empty: {}", result); assert!(
result.contains("empty"),
"Whitespace-only should be empty: {}",
result
);
} }
+12 -3
View File
@@ -22,7 +22,10 @@ fn test_token_accumulation() {
total_tokens: 275, total_tokens: 275,
}; };
window.update_usage_from_response(&usage2); window.update_usage_from_response(&usage2);
assert_eq!(window.used_tokens, 425, "Second call should accumulate to 425 tokens"); assert_eq!(
window.used_tokens, 425,
"Second call should accumulate to 425 tokens"
);
assert_eq!(window.cumulative_tokens, 425, "Cumulative should be 425"); assert_eq!(window.cumulative_tokens, 425, "Cumulative should be 425");
// Third API call with SMALLER token count: 50 prompt + 25 completion = 75 total // Third API call with SMALLER token count: 50 prompt + 25 completion = 75 total
@@ -32,11 +35,17 @@ fn test_token_accumulation() {
total_tokens: 75, total_tokens: 75,
}; };
window.update_usage_from_response(&usage3); window.update_usage_from_response(&usage3);
assert_eq!(window.used_tokens, 500, "Third call should accumulate to 500 tokens"); assert_eq!(
window.used_tokens, 500,
"Third call should accumulate to 500 tokens"
);
assert_eq!(window.cumulative_tokens, 500, "Cumulative should be 500"); assert_eq!(window.cumulative_tokens, 500, "Cumulative should be 500");
// Verify tokens never decrease // Verify tokens never decrease
assert!(window.used_tokens >= 425, "Token count should never decrease!"); assert!(
window.used_tokens >= 425,
"Token count should never decrease!"
);
} }
#[test] #[test]
+34 -11
View File
@@ -1,9 +1,9 @@
use g3_core::{Agent, ToolCall};
use g3_core::ui_writer::UiWriter;
use g3_config::Config; use g3_config::Config;
use g3_core::ui_writer::UiWriter;
use g3_core::{Agent, ToolCall};
use serial_test::serial;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use tempfile::TempDir; use tempfile::TempDir;
use serial_test::serial;
// Mock UI Writer for testing // Mock UI Writer for testing
#[derive(Clone)] #[derive(Clone)]
@@ -47,7 +47,10 @@ impl UiWriter for MockUiWriter {
} }
fn print_system_prompt(&self, _prompt: &str) {} fn print_system_prompt(&self, _prompt: &str) {}
fn print_context_status(&self, message: &str) { fn print_context_status(&self, message: &str) {
self.output.lock().unwrap().push(format!("STATUS: {}", message)); self.output
.lock()
.unwrap()
.push(format!("STATUS: {}", message));
} }
fn print_context_thinning(&self, _message: &str) {} fn print_context_thinning(&self, _message: &str) {}
fn print_tool_header(&self, _tool_name: &str) {} fn print_tool_header(&self, _tool_name: &str) {}
@@ -61,13 +64,21 @@ impl UiWriter for MockUiWriter {
fn print_agent_response(&self, _content: &str) {} fn print_agent_response(&self, _content: &str) {}
fn notify_sse_received(&self) {} fn notify_sse_received(&self) {}
fn flush(&self) {} fn flush(&self) {}
fn wants_full_output(&self) -> bool { false } fn wants_full_output(&self) -> bool {
false
}
fn prompt_user_yes_no(&self, message: &str) -> bool { fn prompt_user_yes_no(&self, message: &str) -> bool {
self.output.lock().unwrap().push(format!("PROMPT: {}", message)); self.output
.lock()
.unwrap()
.push(format!("PROMPT: {}", message));
self.prompt_responses.lock().unwrap().pop().unwrap_or(true) self.prompt_responses.lock().unwrap().pop().unwrap_or(true)
} }
fn prompt_user_choice(&self, message: &str, options: &[&str]) -> usize { fn prompt_user_choice(&self, message: &str, options: &[&str]) -> usize {
self.output.lock().unwrap().push(format!("CHOICE: {} Options: {:?}", message, options)); self.output
.lock()
.unwrap()
.push(format!("CHOICE: {} Options: {:?}", message, options));
self.choice_responses.lock().unwrap().pop().unwrap_or(0) self.choice_responses.lock().unwrap().pop().unwrap_or(0)
} }
} }
@@ -80,7 +91,10 @@ async fn test_todo_staleness_check_matching_sha() {
std::env::set_current_dir(&temp_dir).unwrap(); std::env::set_current_dir(&temp_dir).unwrap();
let sha = "abc123hash"; let sha = "abc123hash";
let content = format!("{{{{Based on the requirements file with SHA256: {}}}}}\n- [ ] Task 1", sha); let content = format!(
"{{{{Based on the requirements file with SHA256: {}}}}}\n- [ ] Task 1",
sha
);
std::fs::write(&todo_path, content).unwrap(); std::fs::write(&todo_path, content).unwrap();
let mut config = Config::default(); let mut config = Config::default();
@@ -109,7 +123,10 @@ async fn test_todo_staleness_check_mismatch_sha_ignore() {
let sha_file = "old_sha"; let sha_file = "old_sha";
let sha_req = "new_sha"; let sha_req = "new_sha";
let content = format!("{{{{Based on the requirements file with SHA256: {}}}}}\n- [ ] Task 1", sha_file); let content = format!(
"{{{{Based on the requirements file with SHA256: {}}}}}\n- [ ] Task 1",
sha_file
);
std::fs::write(&todo_path, content).unwrap(); std::fs::write(&todo_path, content).unwrap();
let mut config = Config::default(); let mut config = Config::default();
@@ -139,7 +156,10 @@ async fn test_todo_staleness_check_mismatch_sha_mark_stale() {
let sha_file = "old_sha"; let sha_file = "old_sha";
let sha_req = "new_sha"; let sha_req = "new_sha";
let content = format!("{{{{Based on the requirements file with SHA256: {}}}}}\n- [ ] Task 1", sha_file); let content = format!(
"{{{{Based on the requirements file with SHA256: {}}}}}\n- [ ] Task 1",
sha_file
);
std::fs::write(&todo_path, content).unwrap(); std::fs::write(&todo_path, content).unwrap();
let mut config = Config::default(); let mut config = Config::default();
@@ -173,7 +193,10 @@ async fn test_todo_staleness_check_disabled() {
let sha_file = "old_sha"; let sha_file = "old_sha";
let sha_req = "new_sha"; let sha_req = "new_sha";
let content = format!("{{{{Based on the requirements file with SHA256: {}}}}}\n- [ ] Task 1", sha_file); let content = format!(
"{{{{Based on the requirements file with SHA256: {}}}}}\n- [ ] Task 1",
sha_file
);
std::fs::write(&todo_path, content).unwrap(); std::fs::write(&todo_path, content).unwrap();
let mut config = Config::default(); let mut config = Config::default();
+91 -36
View File
@@ -43,12 +43,18 @@ impl FlockConfig {
) -> Result<Self> { ) -> Result<Self> {
// Validate project directory // Validate project directory
if !project_dir.exists() { if !project_dir.exists() {
anyhow::bail!("Project directory does not exist: {}", project_dir.display()); anyhow::bail!(
"Project directory does not exist: {}",
project_dir.display()
);
} }
// Check if it's a git repo // Check if it's a git repo
if !project_dir.join(".git").exists() { if !project_dir.join(".git").exists() {
anyhow::bail!("Project directory must be a git repository: {}", project_dir.display()); anyhow::bail!(
"Project directory must be a git repository: {}",
project_dir.display()
);
} }
// Check for flock-requirements.md // Check for flock-requirements.md
@@ -120,10 +126,16 @@ impl FlockMode {
/// Run flock mode /// Run flock mode
pub async fn run(&mut self) -> Result<()> { pub async fn run(&mut self) -> Result<()> {
info!("Starting flock mode with {} segments", self.config.num_segments); info!(
"Starting flock mode with {} segments",
self.config.num_segments
);
// Step 1: Partition requirements // Step 1: Partition requirements
println!("\n🧠 Step 1: Partitioning requirements into {} segments...", self.config.num_segments); println!(
"\n🧠 Step 1: Partitioning requirements into {} segments...",
self.config.num_segments
);
let partitions = self.partition_requirements().await?; let partitions = self.partition_requirements().await?;
// Step 2: Create segment workspaces // Step 2: Create segment workspaces
@@ -131,7 +143,10 @@ impl FlockMode {
self.create_segment_workspaces(&partitions).await?; self.create_segment_workspaces(&partitions).await?;
// Step 3: Run segments in parallel // Step 3: Run segments in parallel
println!("\n🚀 Step 3: Running {} segments in parallel...", self.config.num_segments); println!(
"\n🚀 Step 3: Running {} segments in parallel...",
self.config.num_segments
);
self.run_segments_parallel().await?; self.run_segments_parallel().await?;
// Step 4: Generate final report // Step 4: Generate final report
@@ -226,8 +241,8 @@ impl FlockMode {
.context("Failed to extract partition JSON from agent output")?; .context("Failed to extract partition JSON from agent output")?;
// Parse the partitions // Parse the partitions
let partitions: Vec<serde_json::Value> = serde_json::from_str(&partitions_json) let partitions: Vec<serde_json::Value> =
.context("Failed to parse partition JSON")?; serde_json::from_str(&partitions_json).context("Failed to parse partition JSON")?;
if partitions.len() != self.config.num_segments { if partitions.len() != self.config.num_segments {
warn!( warn!(
@@ -241,9 +256,7 @@ impl FlockMode {
let mut partition_texts = Vec::new(); let mut partition_texts = Vec::new();
for (i, partition) in partitions.iter().enumerate() { for (i, partition) in partitions.iter().enumerate() {
let default_name = format!("module-{}", i + 1); let default_name = format!("module-{}", i + 1);
let module_name = partition["module_name"] let module_name = partition["module_name"].as_str().unwrap_or(&default_name);
.as_str()
.unwrap_or(&default_name);
let requirements = partition["requirements"] let requirements = partition["requirements"]
.as_str() .as_str()
.context("Missing requirements field in partition")?; .context("Missing requirements field in partition")?;
@@ -312,7 +325,9 @@ impl FlockMode {
} }
if candidates.is_empty() { if candidates.is_empty() {
anyhow::bail!("Could not find any partition JSON markers with code fences in agent output"); anyhow::bail!(
"Could not find any partition JSON markers with code fences in agent output"
);
} }
// Try to parse each candidate and return the first valid JSON // Try to parse each candidate and return the first valid JSON
@@ -320,11 +335,20 @@ impl FlockMode {
for (i, candidate) in candidates.iter().enumerate() { for (i, candidate) in candidates.iter().enumerate() {
match serde_json::from_str::<serde_json::Value>(candidate) { match serde_json::from_str::<serde_json::Value>(candidate) {
Ok(_) => { Ok(_) => {
debug!("Successfully parsed JSON from candidate {} of {}", i + 1, candidates.len()); debug!(
"Successfully parsed JSON from candidate {} of {}",
i + 1,
candidates.len()
);
return Ok(candidate.clone()); return Ok(candidate.clone());
} }
Err(e) => { Err(e) => {
debug!("Failed to parse candidate {} of {}: {}", i + 1, candidates.len(), e); debug!(
"Failed to parse candidate {} of {}: {}",
i + 1,
candidates.len(),
e
);
last_error = Some(e); last_error = Some(e);
} }
} }
@@ -349,7 +373,10 @@ impl FlockMode {
for (i, partition) in partitions.iter().enumerate() { for (i, partition) in partitions.iter().enumerate() {
let segment_id = i + 1; let segment_id = i + 1;
let segment_dir = self.config.flock_workspace.join(format!("segment-{}", segment_id)); let segment_dir = self
.config
.flock_workspace
.join(format!("segment-{}", segment_id));
println!(" Creating segment {} workspace...", segment_id); println!(" Creating segment {} workspace...", segment_id);
@@ -360,10 +387,16 @@ impl FlockMode {
// Write segment-requirements.md // Write segment-requirements.md
let requirements_path = segment_dir.join("segment-requirements.md"); let requirements_path = segment_dir.join("segment-requirements.md");
std::fs::write(&requirements_path, partition) std::fs::write(&requirements_path, partition).context(format!(
.context(format!("Failed to write requirements for segment {}", segment_id))?; "Failed to write requirements for segment {}",
segment_id
))?;
println!(" ✓ Segment {} workspace ready at {}", segment_id, segment_dir.display()); println!(
" ✓ Segment {} workspace ready at {}",
segment_id,
segment_dir.display()
);
} }
Ok(()) Ok(())
@@ -393,7 +426,10 @@ impl FlockMode {
let mut handles = Vec::new(); let mut handles = Vec::new();
for segment_id in 1..=self.config.num_segments { for segment_id in 1..=self.config.num_segments {
let segment_dir = self.config.flock_workspace.join(format!("segment-{}", segment_id)); let segment_dir = self
.config
.flock_workspace
.join(format!("segment-{}", segment_id));
let max_turns = self.config.max_turns; let max_turns = self.config.max_turns;
let g3_binary = self.get_g3_binary()?; let g3_binary = self.get_g3_binary()?;
let status_file = self.get_status_file_path(); let status_file = self.get_status_file_path();
@@ -444,10 +480,17 @@ impl FlockMode {
} }
Ok(Err(e)) => { Ok(Err(e)) => {
error!("Segment {} failed: {}", segment_id, e); error!("Segment {} failed: {}", segment_id, e);
let mut segment_status = self.status.segments.get(&segment_id).cloned() let mut segment_status = self
.status
.segments
.get(&segment_id)
.cloned()
.unwrap_or_else(|| SegmentStatus { .unwrap_or_else(|| SegmentStatus {
segment_id, segment_id,
workspace: self.config.flock_workspace.join(format!("segment-{}", segment_id)), workspace: self
.config
.flock_workspace
.join(format!("segment-{}", segment_id)),
state: SegmentState::Failed, state: SegmentState::Failed,
started_at: Utc::now(), started_at: Utc::now(),
completed_at: Some(Utc::now()), completed_at: Some(Utc::now()),
@@ -468,10 +511,17 @@ impl FlockMode {
} }
Err(e) => { Err(e) => {
error!("Segment {} task panicked: {}", segment_id, e); error!("Segment {} task panicked: {}", segment_id, e);
let mut segment_status = self.status.segments.get(&segment_id).cloned() let mut segment_status = self
.status
.segments
.get(&segment_id)
.cloned()
.unwrap_or_else(|| SegmentStatus { .unwrap_or_else(|| SegmentStatus {
segment_id, segment_id,
workspace: self.config.flock_workspace.join(format!("segment-{}", segment_id)), workspace: self
.config
.flock_workspace
.join(format!("segment-{}", segment_id)),
state: SegmentState::Failed, state: SegmentState::Failed,
started_at: Utc::now(), started_at: Utc::now(),
completed_at: Some(Utc::now()), completed_at: Some(Utc::now()),
@@ -527,7 +577,11 @@ async fn run_segment(
status_file: PathBuf, status_file: PathBuf,
session_id: String, session_id: String,
) -> Result<SegmentStatus> { ) -> Result<SegmentStatus> {
info!("Starting segment {} in {}", segment_id, segment_dir.display()); info!(
"Starting segment {} in {}",
segment_id,
segment_dir.display()
);
let mut segment_status = SegmentStatus { let mut segment_status = SegmentStatus {
segment_id, segment_id,
@@ -552,7 +606,9 @@ async fn run_segment(
.arg("--max-turns") .arg("--max-turns")
.arg(max_turns.to_string()) .arg(max_turns.to_string())
.arg("--requirements") .arg("--requirements")
.arg(std::fs::read_to_string(segment_dir.join("segment-requirements.md"))?) .arg(std::fs::read_to_string(
segment_dir.join("segment-requirements.md"),
)?)
.arg("--quiet") // Disable session logging for workers .arg("--quiet") // Disable session logging for workers
.stdout(Stdio::piped()) .stdout(Stdio::piped())
.stderr(Stdio::piped()) .stderr(Stdio::piped())
@@ -615,7 +671,10 @@ async fn run_segment(
} }
// Wait for process to complete // Wait for process to complete
let status = child.wait().await.context("Failed to wait for g3 process")?; let status = child
.wait()
.await
.context("Failed to wait for g3 process")?;
segment_status.completed_at = Some(Utc::now()); segment_status.completed_at = Some(Utc::now());
@@ -636,7 +695,9 @@ async fn run_segment(
let path = entry.path(); let path = entry.path();
if path.extension().and_then(|s| s.to_str()) == Some("json") { if path.extension().and_then(|s| s.to_str()) == Some("json") {
if let Ok(log_content) = std::fs::read_to_string(&path) { if let Ok(log_content) = std::fs::read_to_string(&path) {
if let Ok(log_json) = serde_json::from_str::<serde_json::Value>(&log_content) { if let Ok(log_json) =
serde_json::from_str::<serde_json::Value>(&log_content)
{
// Extract token usage // Extract token usage
if let Some(context) = log_json.get("context_window") { if let Some(context) = log_json.get("context_window") {
if let Some(cumulative) = context.get("cumulative_tokens") { if let Some(cumulative) = context.get("cumulative_tokens") {
@@ -653,8 +714,7 @@ async fn run_segment(
let tool_call_count = messages let tool_call_count = messages
.iter() .iter()
.filter(|msg| { .filter(|msg| {
msg.get("role") msg.get("role").and_then(|r| r.as_str())
.and_then(|r| r.as_str())
== Some("tool") == Some("tool")
}) })
.count(); .count();
@@ -685,12 +745,7 @@ fn update_status_file(
FlockStatus::load_from_file(status_file)? FlockStatus::load_from_file(status_file)?
} else { } else {
// This shouldn't happen, but handle it gracefully // This shouldn't happen, but handle it gracefully
FlockStatus::new( FlockStatus::new(session_id.to_string(), PathBuf::new(), PathBuf::new(), 0)
session_id.to_string(),
PathBuf::new(),
PathBuf::new(),
0,
)
}; };
flock_status.update_segment(segment_status.segment_id, segment_status); flock_status.update_segment(segment_status.segment_id, segment_status);
@@ -896,8 +951,8 @@ The requirements have been partitioned into two logical, largely non-overlapping
.expect("should extract valid JSON from output with multiple markers"); .expect("should extract valid JSON from output with multiple markers");
// Should be able to parse as JSON // Should be able to parse as JSON
let parsed: serde_json::Value = serde_json::from_str(&extracted) let parsed: serde_json::Value =
.expect("extracted content should be valid JSON"); serde_json::from_str(&extracted).expect("extracted content should be valid JSON");
// Verify it's an array with 2 elements // Verify it's an array with 2 elements
assert!(parsed.is_array()); assert!(parsed.is_array());
+40 -10
View File
@@ -185,7 +185,10 @@ impl FlockStatus {
report.push_str(&format!("\n\n🆔 Session ID: {}", self.session_id)); report.push_str(&format!("\n\n🆔 Session ID: {}", self.session_id));
report.push_str(&format!("\n📁 Project: {}", self.project_dir.display())); report.push_str(&format!("\n📁 Project: {}", self.project_dir.display()));
report.push_str(&format!("\n🗂️ Workspace: {}", self.flock_workspace.display())); report.push_str(&format!(
"\n🗂️ Workspace: {}",
self.flock_workspace.display()
));
report.push_str(&format!("\n🔢 Segments: {}", self.num_segments)); report.push_str(&format!("\n🔢 Segments: {}", self.num_segments));
let duration = if let Some(completed) = self.completed_at { let duration = if let Some(completed) = self.completed_at {
@@ -194,20 +197,41 @@ impl FlockStatus {
Utc::now().signed_duration_since(self.started_at) Utc::now().signed_duration_since(self.started_at)
}; };
report.push_str(&format!("\n⏱️ Duration: {:.2}s", duration.num_milliseconds() as f64 / 1000.0)); report.push_str(&format!(
"\n⏱️ Duration: {:.2}s",
duration.num_milliseconds() as f64 / 1000.0
));
// Segment status summary // Segment status summary
report.push_str(&format!("\n\n📈 Segment Status:")); report.push_str(&format!("\n\n📈 Segment Status:"));
report.push_str(&format!("\n • Completed: {}", self.count_by_state(SegmentState::Completed))); report.push_str(&format!(
report.push_str(&format!("\n • Running: {}", self.count_by_state(SegmentState::Running))); "\n • Completed: {}",
report.push_str(&format!("\n • Failed: {}", self.count_by_state(SegmentState::Failed))); self.count_by_state(SegmentState::Completed)
report.push_str(&format!("\n • Pending: {}", self.count_by_state(SegmentState::Pending))); ));
report.push_str(&format!("\n • Cancelled: {}", self.count_by_state(SegmentState::Cancelled))); report.push_str(&format!(
"\n • Running: {}",
self.count_by_state(SegmentState::Running)
));
report.push_str(&format!(
"\n • Failed: {}",
self.count_by_state(SegmentState::Failed)
));
report.push_str(&format!(
"\n • Pending: {}",
self.count_by_state(SegmentState::Pending)
));
report.push_str(&format!(
"\n • Cancelled: {}",
self.count_by_state(SegmentState::Cancelled)
));
// Metrics // Metrics
report.push_str(&format!("\n\n📊 Aggregate Metrics:")); report.push_str(&format!("\n\n📊 Aggregate Metrics:"));
report.push_str(&format!("\n • Total Tokens: {}", self.total_tokens)); report.push_str(&format!("\n • Total Tokens: {}", self.total_tokens));
report.push_str(&format!("\n • Total Tool Calls: {}", self.total_tool_calls)); report.push_str(&format!(
"\n • Total Tool Calls: {}",
self.total_tool_calls
));
report.push_str(&format!("\n • Total Errors: {}", self.total_errors)); report.push_str(&format!("\n • Total Errors: {}", self.total_errors));
// Per-segment details // Per-segment details
@@ -218,11 +242,17 @@ impl FlockStatus {
for (id, segment) in segments { for (id, segment) in segments {
report.push_str(&format!("\n\n Segment {}:", id)); report.push_str(&format!("\n\n Segment {}:", id));
report.push_str(&format!("\n Status: {}", segment.state)); report.push_str(&format!("\n Status: {}", segment.state));
report.push_str(&format!("\n Workspace: {}", segment.workspace.display())); report.push_str(&format!(
"\n Workspace: {}",
segment.workspace.display()
));
report.push_str(&format!("\n Tokens: {}", segment.tokens_used)); report.push_str(&format!("\n Tokens: {}", segment.tokens_used));
report.push_str(&format!("\n Tool Calls: {}", segment.tool_calls)); report.push_str(&format!("\n Tool Calls: {}", segment.tool_calls));
report.push_str(&format!("\n Errors: {}", segment.errors)); report.push_str(&format!("\n Errors: {}", segment.errors));
report.push_str(&format!("\n Turn: {}/{}", segment.current_turn, segment.max_turns)); report.push_str(&format!(
"\n Turn: {}/{}",
segment.current_turn, segment.max_turns
));
if let Some(ref msg) = segment.last_message { if let Some(ref msg) = segment.last_message {
report.push_str(&format!("\n Last Message: {}", msg)); report.push_str(&format!("\n Last Message: {}", msg));
+1 -2
View File
@@ -283,8 +283,7 @@ mod tests {
assert!(json.contains("Completed")); assert!(json.contains("Completed"));
// Deserialize back // Deserialize back
let deserialized: FlockStatus = let deserialized: FlockStatus = serde_json::from_str(&json).expect("Failed to deserialize");
serde_json::from_str(&json).expect("Failed to deserialize");
assert_eq!(deserialized.session_id, "test-session"); assert_eq!(deserialized.session_id, "test-session");
assert_eq!(deserialized.segments.len(), 1); assert_eq!(deserialized.segments.len(), 1);
assert_eq!(deserialized.total_tokens, 1000); assert_eq!(deserialized.total_tokens, 1000);
@@ -71,7 +71,7 @@ fn create_test_project(name: &str) -> TempDir {
} }
#[test] #[test]
fn test_flock_config_validation() { fn test_flock_config_validation() {
let temp_dir = TempDir::new().unwrap(); let temp_dir = TempDir::new().unwrap();
let project_path = temp_dir.path().to_path_buf(); let project_path = temp_dir.path().to_path_buf();
let workspace_path = temp_dir.path().join("workspace"); let workspace_path = temp_dir.path().join("workspace");
@@ -213,8 +213,7 @@ fn test_multiple_segment_clones() {
assert!(segment2.exists()); assert!(segment2.exists());
// Modify segment 1 // Modify segment 1
fs::write(segment1.join("test.txt"), "segment 1") fs::write(segment1.join("test.txt"), "segment 1").expect("Failed to write to segment 1");
.expect("Failed to write to segment 1");
// Verify segment 2 is unaffected // Verify segment 2 is unaffected
assert!(!segment2.join("test.txt").exists()); assert!(!segment2.join("test.txt").exists());
@@ -236,8 +235,11 @@ fn test_segment_requirements_creation() {
// Create segment-requirements.md (what flock mode does) // Create segment-requirements.md (what flock mode does)
let segment_requirements = "# Module A\n\nImplement module A functionality\n"; let segment_requirements = "# Module A\n\nImplement module A functionality\n";
fs::write(segment_dir.join("segment-requirements.md"), segment_requirements) fs::write(
.expect("Failed to write segment requirements"); segment_dir.join("segment-requirements.md"),
segment_requirements,
)
.expect("Failed to write segment requirements");
// Verify it was created // Verify it was created
assert!(segment_dir.join("segment-requirements.md").exists()); assert!(segment_dir.join("segment-requirements.md").exists());
+60 -32
View File
@@ -1,9 +1,9 @@
use anyhow::Result; use anyhow::Result;
use regex::Regex; use regex::Regex;
use std::io::Write;
use std::process::Command; use std::process::Command;
use tempfile::NamedTempFile; use tempfile::NamedTempFile;
use std::io::Write; use tracing::{debug, error, info};
use tracing::{info, debug, error};
/// Expand tilde (~) in a path to the user's home directory /// Expand tilde (~) in a path to the user's home directory
fn expand_tilde(path: &str) -> String { fn expand_tilde(path: &str) -> String {
@@ -35,17 +35,29 @@ impl CodeExecutor {
/// Extract code blocks from LLM response and execute them /// Extract code blocks from LLM response and execute them
pub async fn execute_from_response(&self, response: &str) -> Result<String> { pub async fn execute_from_response(&self, response: &str) -> Result<String> {
self.execute_from_response_with_options(response, true).await self.execute_from_response_with_options(response, true)
.await
} }
/// Extract code blocks from LLM response and execute them with UI options /// Extract code blocks from LLM response and execute them with UI options
pub async fn execute_from_response_with_options(&self, response: &str, show_code: bool) -> Result<String> { pub async fn execute_from_response_with_options(
debug!("CodeExecutor received response ({} chars): {}", response.len(), response); &self,
response: &str,
show_code: bool,
) -> Result<String> {
debug!(
"CodeExecutor received response ({} chars): {}",
response.len(),
response
);
let code_blocks = self.extract_code_blocks(response)?; let code_blocks = self.extract_code_blocks(response)?;
if code_blocks.is_empty() { if code_blocks.is_empty() {
if show_code { if show_code {
return Ok(format!("⚠️ No executable code blocks found in response.\n\n{}", response)); return Ok(format!(
"⚠️ No executable code blocks found in response.\n\n{}",
response
));
} else { } else {
return Ok("⚠️ No executable code found.".to_string()); return Ok("⚠️ No executable code found.".to_string());
} }
@@ -108,12 +120,16 @@ impl CodeExecutor {
// Pattern 1: Standard markdown format ```language\ncode``` // Pattern 1: Standard markdown format ```language\ncode```
let markdown_re = Regex::new(r"(?s)```(\w+)?\n(.*?)```")?; let markdown_re = Regex::new(r"(?s)```(\w+)?\n(.*?)```")?;
for cap in markdown_re.captures_iter(text) { for cap in markdown_re.captures_iter(text) {
let language = cap.get(1) let language = cap
.get(1)
.map(|m| m.as_str().to_lowercase()) .map(|m| m.as_str().to_lowercase())
.unwrap_or_else(|| "bash".to_string()); // Default to bash .unwrap_or_else(|| "bash".to_string()); // Default to bash
let code = cap.get(2).map(|m| m.as_str()).unwrap_or("").trim(); let code = cap.get(2).map(|m| m.as_str()).unwrap_or("").trim();
debug!("Found markdown code block - language: '{}', code: '{}'", language, code); debug!(
"Found markdown code block - language: '{}', code: '{}'",
language, code
);
if !code.is_empty() { if !code.is_empty() {
blocks.push((language, code.to_string())); blocks.push((language, code.to_string()));
@@ -131,7 +147,10 @@ impl CodeExecutor {
let language = open_lang.to_lowercase(); let language = open_lang.to_lowercase();
let code = cap.get(2).map(|m| m.as_str()).unwrap_or("").trim(); let code = cap.get(2).map(|m| m.as_str()).unwrap_or("").trim();
debug!("Found bracket code block - language: '{}', code: '{}'", language, code); debug!(
"Found bracket code block - language: '{}', code: '{}'",
language, code
);
if !code.is_empty() { if !code.is_empty() {
blocks.push((language, code.to_string())); blocks.push((language, code.to_string()));
@@ -163,9 +182,7 @@ impl CodeExecutor {
temp_file.write_all(code.as_bytes())?; temp_file.write_all(code.as_bytes())?;
let temp_path = temp_file.path(); let temp_path = temp_file.path();
let output = Command::new("python3") let output = Command::new("python3").arg(temp_path).output()?;
.arg(temp_path)
.output()?;
Ok(ExecutionResult { Ok(ExecutionResult {
stdout: String::from_utf8_lossy(&output.stdout).to_string(), stdout: String::from_utf8_lossy(&output.stdout).to_string(),
@@ -202,10 +219,7 @@ impl CodeExecutor {
}); });
} }
let output = Command::new("bash") let output = Command::new("bash").arg("-c").arg(code).output()?;
.arg("-c")
.arg(code)
.output()?;
Ok(ExecutionResult { Ok(ExecutionResult {
stdout: String::from_utf8_lossy(&output.stdout).to_string(), stdout: String::from_utf8_lossy(&output.stdout).to_string(),
@@ -221,9 +235,7 @@ impl CodeExecutor {
temp_file.write_all(code.as_bytes())?; temp_file.write_all(code.as_bytes())?;
let temp_path = temp_file.path(); let temp_path = temp_file.path();
let output = Command::new("node") let output = Command::new("node").arg(temp_path).output()?;
.arg(temp_path)
.output()?;
Ok(ExecutionResult { Ok(ExecutionResult {
stdout: String::from_utf8_lossy(&output.stdout).to_string(), stdout: String::from_utf8_lossy(&output.stdout).to_string(),
@@ -251,9 +263,10 @@ impl CodeExecutor {
pub async fn execute_bash_streaming<R: OutputReceiver>( pub async fn execute_bash_streaming<R: OutputReceiver>(
&self, &self,
code: &str, code: &str,
receiver: &R receiver: &R,
) -> Result<ExecutionResult> { ) -> Result<ExecutionResult> {
self.execute_bash_streaming_in_dir(code, receiver, None).await self.execute_bash_streaming_in_dir(code, receiver, None)
.await
} }
/// Execute bash command with streaming output in a specific directory /// Execute bash command with streaming output in a specific directory
@@ -271,13 +284,25 @@ impl CodeExecutor {
debug!("========== execute_bash_streaming_in_dir START =========="); debug!("========== execute_bash_streaming_in_dir START ==========");
debug!("Code to execute: {}", code); debug!("Code to execute: {}", code);
debug!("Working directory parameter: {:?}", working_dir); debug!("Working directory parameter: {:?}", working_dir);
debug!("FULL DIAGNOSTIC: code='{}', working_dir={:?}", code, working_dir); debug!(
"FULL DIAGNOSTIC: code='{}', working_dir={:?}",
code, working_dir
);
if let Some(dir) = working_dir { if let Some(dir) = working_dir {
debug!("Working dir exists check: {}", std::path::Path::new(dir).exists()); debug!(
debug!("Working dir is_dir check: {}", std::path::Path::new(dir).is_dir()); "Working dir exists check: {}",
std::path::Path::new(dir).exists()
);
debug!(
"Working dir is_dir check: {}",
std::path::Path::new(dir).is_dir()
);
} }
debug!("Current process working directory: {:?}", std::env::current_dir()); debug!(
"Current process working directory: {:?}",
std::env::current_dir()
);
// Check if this is a detached/daemon command that should run independently // Check if this is a detached/daemon command that should run independently
// Look for patterns like: setsid, nohup with &, or explicit backgrounding with disown // Look for patterns like: setsid, nohup with &, or explicit backgrounding with disown
@@ -289,8 +314,7 @@ impl CodeExecutor {
if is_detached { if is_detached {
// For detached commands, just spawn and return immediately // For detached commands, just spawn and return immediately
let mut cmd = TokioCommand::new("bash"); let mut cmd = TokioCommand::new("bash");
cmd.arg("-c") cmd.arg("-c").arg(code);
.arg(code);
// Set working directory if provided // Set working directory if provided
if let Some(dir) = working_dir { if let Some(dir) = working_dir {
@@ -320,8 +344,14 @@ impl CodeExecutor {
debug!("Setting current_dir on command to: {}", dir); debug!("Setting current_dir on command to: {}", dir);
let expanded_dir = expand_tilde(dir); let expanded_dir = expand_tilde(dir);
debug!("Expanded working dir: {}", expanded_dir); debug!("Expanded working dir: {}", expanded_dir);
debug!("Expanded dir exists: {}", std::path::Path::new(&expanded_dir).exists()); debug!(
debug!("Expanded dir is_dir: {}", std::path::Path::new(&expanded_dir).is_dir()); "Expanded dir exists: {}",
std::path::Path::new(&expanded_dir).exists()
);
debug!(
"Expanded dir is_dir: {}",
std::path::Path::new(&expanded_dir).is_dir()
);
cmd.current_dir(&expanded_dir); cmd.current_dir(&expanded_dir);
} }
@@ -418,9 +448,7 @@ pub fn is_llvm_tools_installed() -> Result<bool> {
/// Check if cargo-llvm-cov is installed /// Check if cargo-llvm-cov is installed
pub fn is_cargo_llvm_cov_installed() -> Result<bool> { pub fn is_cargo_llvm_cov_installed() -> Result<bool> {
let output = Command::new("cargo") let output = Command::new("cargo").args(&["--list"]).output()?;
.args(&["--list"])
.output()?;
let installed = String::from_utf8_lossy(&output.stdout) let installed = String::from_utf8_lossy(&output.stdout)
.lines() .lines()
+4 -1
View File
@@ -291,7 +291,10 @@ pub fn explore_kotlin(path: &str) -> String {
// Build files // Build files
report.push_str("--- Build Configuration ---\n"); report.push_str("--- Build Configuration ---\n");
let build = run_command("cat build.gradle.kts 2>/dev/null | head -50 || cat build.gradle 2>/dev/null | head -50", path); let build = run_command(
"cat build.gradle.kts 2>/dev/null | head -50 || cat build.gradle 2>/dev/null | head -50",
path,
);
report.push_str(&build); report.push_str(&build);
report.push('\n'); report.push('\n');
+6 -3
View File
@@ -9,11 +9,11 @@ pub mod prompts;
pub use code_explore::explore_codebase; pub use code_explore::explore_codebase;
use anyhow::Result; use anyhow::Result;
use g3_providers::{CompletionRequest, LLMProvider, Message, MessageRole};
use chrono::Local; use chrono::Local;
use g3_providers::{CompletionRequest, LLMProvider, Message, MessageRole};
use prompts::{DISCOVERY_REQUIREMENTS_PROMPT, DISCOVERY_SYSTEM_PROMPT};
use std::fs::{self, OpenOptions}; use std::fs::{self, OpenOptions};
use std::io::Write; use std::io::Write;
use prompts::{DISCOVERY_REQUIREMENTS_PROMPT, DISCOVERY_SYSTEM_PROMPT};
/// Type alias for a status callback function /// Type alias for a status callback function
pub type StatusCallback = Box<dyn Fn(&str) + Send + Sync>; pub type StatusCallback = Box<dyn Fn(&str) + Send + Sync>;
@@ -94,7 +94,10 @@ pub async fn get_initial_discovery_messages(
// Step 5: Extract shell commands from the response // Step 5: Extract shell commands from the response
let shell_commands = extract_shell_commands(&response.content); let shell_commands = extract_shell_commands(&response.content);
status(&format!("📋 Extracted {} discovery commands", shell_commands.len())); status(&format!(
"📋 Extracted {} discovery commands",
shell_commands.len()
));
// Write the discovery commands to logs directory // Write the discovery commands to logs directory
write_discovery_commands(&shell_commands)?; write_discovery_commands(&shell_commands)?;
+4 -2
View File
@@ -31,12 +31,14 @@ fn test_log_files_created() {
// Test writing discovery commands // Test writing discovery commands
let commands_filename = format!("logs/discovery_commands_{}.log", timestamp); let commands_filename = format!("logs/discovery_commands_{}.log", timestamp);
let test_commands = "# Discovery Commands\n# Generated by g3-planner\n\nls -la\ncat README.md\n"; let test_commands =
"# Discovery Commands\n# Generated by g3-planner\n\nls -la\ncat README.md\n";
fs::write(&commands_filename, test_commands).expect("Failed to write discovery commands"); fs::write(&commands_filename, test_commands).expect("Failed to write discovery commands");
assert!(Path::new(&commands_filename).exists()); assert!(Path::new(&commands_filename).exists());
let content = fs::read_to_string(&commands_filename).expect("Failed to read discovery commands"); let content =
fs::read_to_string(&commands_filename).expect("Failed to read discovery commands");
assert_eq!(content, test_commands); assert_eq!(content, test_commands);
// Clean up // Clean up
+165 -85
View File
@@ -188,12 +188,17 @@ impl AnthropicProvider {
}; };
// Extract properties and required fields from the input schema // Extract properties and required fields from the input schema
if let Ok(schema_obj) = serde_json::from_value::<serde_json::Map<String, serde_json::Value>>(tool.input_schema.clone()) { if let Ok(schema_obj) = serde_json::from_value::<
serde_json::Map<String, serde_json::Value>,
>(tool.input_schema.clone())
{
if let Some(properties) = schema_obj.get("properties") { if let Some(properties) = schema_obj.get("properties") {
schema.properties = properties.clone(); schema.properties = properties.clone();
} }
if let Some(required) = schema_obj.get("required") { if let Some(required) = schema_obj.get("required") {
if let Ok(required_vec) = serde_json::from_value::<Vec<String>>(required.clone()) { if let Ok(required_vec) =
serde_json::from_value::<Vec<String>>(required.clone())
{
schema.required = Some(required_vec); schema.required = Some(required_vec);
} }
} }
@@ -208,7 +213,10 @@ impl AnthropicProvider {
.collect() .collect()
} }
fn convert_messages(&self, messages: &[Message]) -> Result<(Option<String>, Vec<AnthropicMessage>)> { fn convert_messages(
&self,
messages: &[Message],
) -> Result<(Option<String>, Vec<AnthropicMessage>)> {
let mut system_message = None; let mut system_message = None;
let mut anthropic_messages = Vec::new(); let mut anthropic_messages = Vec::new();
@@ -225,7 +233,9 @@ impl AnthropicProvider {
role: "user".to_string(), role: "user".to_string(),
content: vec![AnthropicContent::Text { content: vec![AnthropicContent::Text {
text: message.content.clone(), text: message.content.clone(),
cache_control: message.cache_control.as_ref() cache_control: message
.cache_control
.as_ref()
.map(Self::convert_cache_control), .map(Self::convert_cache_control),
}], }],
}); });
@@ -235,7 +245,9 @@ impl AnthropicProvider {
role: "assistant".to_string(), role: "assistant".to_string(),
content: vec![AnthropicContent::Text { content: vec![AnthropicContent::Text {
text: message.content.clone(), text: message.content.clone(),
cache_control: message.cache_control.as_ref() cache_control: message
.cache_control
.as_ref()
.map(Self::convert_cache_control), .map(Self::convert_cache_control),
}], }],
}); });
@@ -257,7 +269,9 @@ impl AnthropicProvider {
let (system, anthropic_messages) = self.convert_messages(messages)?; let (system, anthropic_messages) = self.convert_messages(messages)?;
if anthropic_messages.is_empty() { if anthropic_messages.is_empty() {
return Err(anyhow!("At least one user or assistant message is required")); return Err(anyhow!(
"At least one user or assistant message is required"
));
} }
// Convert tools if provided // Convert tools if provided
@@ -312,7 +326,8 @@ impl AnthropicProvider {
let valid_up_to = e.valid_up_to(); let valid_up_to = e.valid_up_to();
if valid_up_to > 0 { if valid_up_to > 0 {
// We have some valid UTF-8, extract it and keep the rest for next iteration // We have some valid UTF-8, extract it and keep the rest for next iteration
let valid_bytes = byte_buffer.drain(..valid_up_to).collect::<Vec<_>>(); let valid_bytes =
byte_buffer.drain(..valid_up_to).collect::<Vec<_>>();
std::str::from_utf8(&valid_bytes).unwrap().to_string() std::str::from_utf8(&valid_bytes).unwrap().to_string()
} else { } else {
// No valid UTF-8 at all, skip this chunk and continue // No valid UTF-8 at all, skip this chunk and continue
@@ -346,7 +361,11 @@ impl AnthropicProvider {
content: String::new(), content: String::new(),
finished: true, finished: true,
usage: accumulated_usage.clone(), usage: accumulated_usage.clone(),
tool_calls: if current_tool_calls.is_empty() { None } else { Some(current_tool_calls.clone()) }, tool_calls: if current_tool_calls.is_empty() {
None
} else {
Some(current_tool_calls.clone())
},
}; };
if tx.send(Ok(final_chunk)).await.is_err() { if tx.send(Ok(final_chunk)).await.is_err() {
debug!("Receiver dropped, stopping stream"); debug!("Receiver dropped, stopping stream");
@@ -358,7 +377,10 @@ impl AnthropicProvider {
match serde_json::from_str::<AnthropicStreamEvent>(data) { match serde_json::from_str::<AnthropicStreamEvent>(data) {
Ok(event) => { Ok(event) => {
debug!("Parsed event type: {}, event: {:?}", event.event_type, event); debug!(
"Parsed event type: {}, event: {:?}",
event.event_type, event
);
match event.event_type.as_str() { match event.event_type.as_str() {
"message_start" => { "message_start" => {
// Extract usage data from message_start event // Extract usage data from message_start event
@@ -367,17 +389,28 @@ impl AnthropicProvider {
accumulated_usage = Some(Usage { accumulated_usage = Some(Usage {
prompt_tokens: usage.input_tokens, prompt_tokens: usage.input_tokens,
completion_tokens: usage.output_tokens, completion_tokens: usage.output_tokens,
total_tokens: usage.input_tokens + usage.output_tokens, total_tokens: usage.input_tokens
+ usage.output_tokens,
}); });
debug!("Captured usage from message_start: {:?}", accumulated_usage); debug!(
"Captured usage from message_start: {:?}",
accumulated_usage
);
} }
} }
} }
"content_block_start" => { "content_block_start" => {
debug!("Received content_block_start event: {:?}", event); debug!(
"Received content_block_start event: {:?}",
event
);
if let Some(content_block) = event.content_block { if let Some(content_block) = event.content_block {
match content_block { match content_block {
AnthropicContent::ToolUse { id, name, input } => { AnthropicContent::ToolUse {
id,
name,
input,
} => {
debug!("Found tool use in content_block_start: id={}, name={}, input={:?}", id, name, input); debug!("Found tool use in content_block_start: id={}, name={}, input={:?}", id, name, input);
// For native tool calls, create the tool call immediately if we have complete args // For native tool calls, create the tool call immediately if we have complete args
@@ -389,7 +422,12 @@ impl AnthropicProvider {
}; };
// Check if we already have complete arguments // Check if we already have complete arguments
if !input.is_null() && input != serde_json::Value::Object(serde_json::Map::new()) { if !input.is_null()
&& input
!= serde_json::Value::Object(
serde_json::Map::new(),
)
{
// We have complete arguments, send the tool call immediately // We have complete arguments, send the tool call immediately
debug!("Tool call has complete args, sending immediately: {:?}", tool_call); debug!("Tool call has complete args, sending immediately: {:?}", tool_call);
let chunk = CompletionChunk { let chunk = CompletionChunk {
@@ -410,7 +448,10 @@ impl AnthropicProvider {
} }
} }
_ => { _ => {
debug!("Non-tool content block: {:?}", content_block); debug!(
"Non-tool content block: {:?}",
content_block
);
} }
} }
} }
@@ -418,7 +459,11 @@ impl AnthropicProvider {
"content_block_delta" => { "content_block_delta" => {
if let Some(delta) = event.delta { if let Some(delta) = event.delta {
if let Some(text) = delta.text { if let Some(text) = delta.text {
debug!("Sending text chunk of length {}: '{}'", text.len(), text); debug!(
"Sending text chunk of length {}: '{}'",
text.len(),
text
);
let chunk = CompletionChunk { let chunk = CompletionChunk {
content: text, content: text,
finished: false, finished: false,
@@ -432,25 +477,45 @@ impl AnthropicProvider {
} }
// Handle partial JSON for tool calls // Handle partial JSON for tool calls
if let Some(partial_json) = delta.partial_json { if let Some(partial_json) = delta.partial_json {
debug!("Received partial JSON: {}", partial_json); debug!(
"Received partial JSON: {}",
partial_json
);
partial_tool_json.push_str(&partial_json); partial_tool_json.push_str(&partial_json);
debug!("Accumulated tool JSON: {}", partial_tool_json); debug!(
"Accumulated tool JSON: {}",
partial_tool_json
);
} }
} }
} }
"content_block_stop" => { "content_block_stop" => {
// Tool call block is complete - now parse the accumulated JSON // Tool call block is complete - now parse the accumulated JSON
if !current_tool_calls.is_empty() && !partial_tool_json.is_empty() { if !current_tool_calls.is_empty()
debug!("Parsing complete tool JSON: {}", partial_tool_json); && !partial_tool_json.is_empty()
{
debug!(
"Parsing complete tool JSON: {}",
partial_tool_json
);
// Parse the accumulated JSON and update the last tool call // Parse the accumulated JSON and update the last tool call
if let Ok(parsed_args) = serde_json::from_str::<serde_json::Value>(&partial_tool_json) { if let Ok(parsed_args) =
if let Some(last_tool) = current_tool_calls.last_mut() { serde_json::from_str::<serde_json::Value>(
&partial_tool_json,
)
{
if let Some(last_tool) =
current_tool_calls.last_mut()
{
last_tool.args = parsed_args; last_tool.args = parsed_args;
debug!("Updated tool call with complete args: {:?}", last_tool); debug!("Updated tool call with complete args: {:?}", last_tool);
} }
} else { } else {
debug!("Failed to parse accumulated JSON: {}", partial_tool_json); debug!(
"Failed to parse accumulated JSON: {}",
partial_tool_json
);
} }
// Clear the accumulator // Clear the accumulator
@@ -478,7 +543,11 @@ impl AnthropicProvider {
content: String::new(), content: String::new(),
finished: true, finished: true,
usage: accumulated_usage.clone(), usage: accumulated_usage.clone(),
tool_calls: if current_tool_calls.is_empty() { None } else { Some(current_tool_calls.clone()) }, tool_calls: if current_tool_calls.is_empty() {
None
} else {
Some(current_tool_calls.clone())
},
}; };
if tx.send(Ok(final_chunk)).await.is_err() { if tx.send(Ok(final_chunk)).await.is_err() {
debug!("Receiver dropped, stopping stream"); debug!("Receiver dropped, stopping stream");
@@ -490,7 +559,10 @@ impl AnthropicProvider {
if let Some(error) = event.error { if let Some(error) = event.error {
error!("Anthropic API error: {:?}", error); error!("Anthropic API error: {:?}", error);
let _ = tx let _ = tx
.send(Err(anyhow!("Anthropic API error: {:?}", error))) .send(Err(anyhow!(
"Anthropic API error: {:?}",
error
)))
.await; .await;
break; // Break to let stream exhaust naturally break; // Break to let stream exhaust naturally
} }
@@ -524,7 +596,11 @@ impl AnthropicProvider {
content: String::new(), content: String::new(),
finished: true, finished: true,
usage: accumulated_usage.clone(), usage: accumulated_usage.clone(),
tool_calls: if current_tool_calls.is_empty() { None } else { Some(current_tool_calls) }, tool_calls: if current_tool_calls.is_empty() {
None
} else {
Some(current_tool_calls)
},
}; };
let _ = tx.send(Ok(final_chunk)).await; let _ = tx.send(Ok(final_chunk)).await;
accumulated_usage accumulated_usage
@@ -547,11 +623,13 @@ impl LLMProvider for AnthropicProvider {
request.tools.as_deref(), request.tools.as_deref(),
false, false,
max_tokens, max_tokens,
temperature temperature,
)?; )?;
debug!("Sending request to Anthropic API: model={}, max_tokens={}, temperature={}", debug!(
request_body.model, request_body.max_tokens, request_body.temperature); "Sending request to Anthropic API: model={}, max_tokens={}, temperature={}",
request_body.model, request_body.max_tokens, request_body.temperature
);
let response = self let response = self
.create_request_builder(false) .create_request_builder(false)
@@ -588,7 +666,8 @@ impl LLMProvider for AnthropicProvider {
let usage = Usage { let usage = Usage {
prompt_tokens: anthropic_response.usage.input_tokens, prompt_tokens: anthropic_response.usage.input_tokens,
completion_tokens: anthropic_response.usage.output_tokens, completion_tokens: anthropic_response.usage.output_tokens,
total_tokens: anthropic_response.usage.input_tokens + anthropic_response.usage.output_tokens, total_tokens: anthropic_response.usage.input_tokens
+ anthropic_response.usage.output_tokens,
}; };
debug!( debug!(
@@ -617,14 +696,20 @@ impl LLMProvider for AnthropicProvider {
request.tools.as_deref(), request.tools.as_deref(),
true, true,
max_tokens, max_tokens,
temperature temperature,
)?; )?;
debug!("Sending streaming request to Anthropic API: model={}, max_tokens={}, temperature={}", debug!(
request_body.model, request_body.max_tokens, request_body.temperature); "Sending streaming request to Anthropic API: model={}, max_tokens={}, temperature={}",
request_body.model, request_body.max_tokens, request_body.temperature
);
// Debug: Log the full request body // Debug: Log the full request body
debug!("Full request body: {}", serde_json::to_string_pretty(&request_body).unwrap_or_else(|_| "Failed to serialize".to_string())); debug!(
"Full request body: {}",
serde_json::to_string_pretty(&request_body)
.unwrap_or_else(|_| "Failed to serialize".to_string())
);
let response = self let response = self
.create_request_builder(true) .create_request_builder(true)
@@ -798,17 +883,14 @@ mod tests {
#[test] #[test]
fn test_message_conversion() { fn test_message_conversion() {
let provider = AnthropicProvider::new( let provider =
"test-key".to_string(), AnthropicProvider::new("test-key".to_string(), None, None, None, None, None).unwrap();
None,
None,
None,
None,
None,
).unwrap();
let messages = vec![ let messages = vec![
Message::new(MessageRole::System, "You are a helpful assistant.".to_string()), Message::new(
MessageRole::System,
"You are a helpful assistant.".to_string(),
),
Message::new(MessageRole::User, "Hello!".to_string()), Message::new(MessageRole::User, "Hello!".to_string()),
Message::new(MessageRole::Assistant, "Hi there!".to_string()), Message::new(MessageRole::Assistant, "Hi there!".to_string()),
]; ];
@@ -830,7 +912,8 @@ mod tests {
Some(0.5), Some(0.5),
None, None,
None, None,
).unwrap(); )
.unwrap();
let messages = vec![Message::new(MessageRole::User, "Test message".to_string())]; let messages = vec![Message::new(MessageRole::User, "Test message".to_string())];
@@ -848,31 +931,23 @@ mod tests {
#[test] #[test]
fn test_tool_conversion() { fn test_tool_conversion() {
let provider = AnthropicProvider::new( let provider =
"test-key".to_string(), AnthropicProvider::new("test-key".to_string(), None, None, None, None, None).unwrap();
None,
None,
None,
None,
None,
).unwrap();
let tools = vec![ let tools = vec![Tool {
Tool { name: "get_weather".to_string(),
name: "get_weather".to_string(), description: "Get the current weather".to_string(),
description: "Get the current weather".to_string(), input_schema: serde_json::json!({
input_schema: serde_json::json!({ "type": "object",
"type": "object", "properties": {
"properties": { "location": {
"location": { "type": "string",
"type": "string", "description": "The city and state"
"description": "The city and state" }
} },
}, "required": ["location"]
"required": ["location"] }),
}), }];
},
];
let anthropic_tools = provider.convert_tools(&tools); let anthropic_tools = provider.convert_tools(&tools);
@@ -881,19 +956,16 @@ mod tests {
assert_eq!(anthropic_tools[0].description, "Get the current weather"); assert_eq!(anthropic_tools[0].description, "Get the current weather");
assert_eq!(anthropic_tools[0].input_schema.schema_type, "object"); assert_eq!(anthropic_tools[0].input_schema.schema_type, "object");
assert!(anthropic_tools[0].input_schema.required.is_some()); assert!(anthropic_tools[0].input_schema.required.is_some());
assert_eq!(anthropic_tools[0].input_schema.required.as_ref().unwrap()[0], "location"); assert_eq!(
anthropic_tools[0].input_schema.required.as_ref().unwrap()[0],
"location"
);
} }
#[test] #[test]
fn test_cache_control_serialization() { fn test_cache_control_serialization() {
let provider = AnthropicProvider::new( let provider =
"test-key".to_string(), AnthropicProvider::new("test-key".to_string(), None, None, None, None, None).unwrap();
None,
None,
None,
None,
None,
).unwrap();
// Test message WITHOUT cache_control // Test message WITHOUT cache_control
let messages_without = vec![Message::new(MessageRole::User, "Hello".to_string())]; let messages_without = vec![Message::new(MessageRole::User, "Hello".to_string())];
@@ -904,8 +976,10 @@ mod tests {
// Check if cache_control appears in the JSON // Check if cache_control appears in the JSON
if json_without.contains("cache_control") { if json_without.contains("cache_control") {
println!("WARNING: JSON contains 'cache_control' field when not configured!"); println!("WARNING: JSON contains 'cache_control' field when not configured!");
assert!(!json_without.contains("\"cache_control\":null"), assert!(
"JSON should not contain 'cache_control: null'"); !json_without.contains("\"cache_control\":null"),
"JSON should not contain 'cache_control: null'"
);
} }
// Test message WITH cache_control // Test message WITH cache_control
@@ -918,13 +992,19 @@ mod tests {
let json_with = serde_json::to_string(&anthropic_messages_with).unwrap(); let json_with = serde_json::to_string(&anthropic_messages_with).unwrap();
println!("Anthropic JSON with cache_control: {}", json_with); println!("Anthropic JSON with cache_control: {}", json_with);
assert!(json_with.contains("cache_control"), assert!(
"JSON should contain 'cache_control' field when configured"); json_with.contains("cache_control"),
assert!(json_with.contains("ephemeral"), "JSON should contain 'cache_control' field when configured"
"JSON should contain 'ephemeral' type"); );
assert!(
json_with.contains("ephemeral"),
"JSON should contain 'ephemeral' type"
);
// The key assertion: when cache_control is None, it should not appear in JSON // The key assertion: when cache_control is None, it should not appear in JSON
assert!(!json_without.contains("cache_control") || !json_without.contains("null"), assert!(
"JSON should not contain 'cache_control' field or null values when not configured"); !json_without.contains("cache_control") || !json_without.contains("null"),
"JSON should not contain 'cache_control' field or null values when not configured"
);
} }
} }
+28 -10
View File
@@ -326,7 +326,8 @@ impl DatabricksProvider {
let valid_up_to = e.valid_up_to(); let valid_up_to = e.valid_up_to();
if valid_up_to > 0 { if valid_up_to > 0 {
// We have some valid UTF-8, extract it and keep the rest for next iteration // We have some valid UTF-8, extract it and keep the rest for next iteration
let valid_bytes = byte_buffer.drain(..valid_up_to).collect::<Vec<_>>(); let valid_bytes =
byte_buffer.drain(..valid_up_to).collect::<Vec<_>>();
std::str::from_utf8(&valid_bytes).unwrap().to_string() std::str::from_utf8(&valid_bytes).unwrap().to_string()
} else { } else {
// No valid UTF-8 at all, skip this chunk and continue // No valid UTF-8 at all, skip this chunk and continue
@@ -610,8 +611,12 @@ impl DatabricksProvider {
// Log final state // Log final state
debug!("Stream ended after {} chunks", chunk_count); debug!("Stream ended after {} chunks", chunk_count);
debug!("Final state: buffer_len={}, incomplete_data_line_len={}, byte_buffer_len={}", debug!(
buffer.len(), incomplete_data_line.len(), byte_buffer.len()); "Final state: buffer_len={}, incomplete_data_line_len={}, byte_buffer_len={}",
buffer.len(),
incomplete_data_line.len(),
byte_buffer.len()
);
debug!("Accumulated tool calls: {}", current_tool_calls.len()); debug!("Accumulated tool calls: {}", current_tool_calls.len());
// If we have any remaining data in buffers, log it for debugging // If we have any remaining data in buffers, log it for debugging
@@ -1181,7 +1186,10 @@ mod tests {
.unwrap(); .unwrap();
let messages = vec![ let messages = vec![
Message::new(MessageRole::System, "You are a helpful assistant.".to_string()), Message::new(
MessageRole::System,
"You are a helpful assistant.".to_string(),
),
Message::new(MessageRole::User, "Hello!".to_string()), Message::new(MessageRole::User, "Hello!".to_string()),
Message::new(MessageRole::Assistant, "Hi there!".to_string()), Message::new(MessageRole::Assistant, "Hi there!".to_string()),
]; ];
@@ -1306,8 +1314,10 @@ mod tests {
let json_without = serde_json::to_string(&databricks_messages_without).unwrap(); let json_without = serde_json::to_string(&databricks_messages_without).unwrap();
println!("JSON without cache_control: {}", json_without); println!("JSON without cache_control: {}", json_without);
assert!(!json_without.contains("cache_control"), assert!(
"JSON should not contain 'cache_control' field when not configured"); !json_without.contains("cache_control"),
"JSON should not contain 'cache_control' field when not configured"
);
// Test message WITH cache_control - should still NOT include it (Databricks doesn't support it) // Test message WITH cache_control - should still NOT include it (Databricks doesn't support it)
let messages_with = vec![Message::with_cache_control( let messages_with = vec![Message::with_cache_control(
@@ -1319,8 +1329,10 @@ mod tests {
let json_with = serde_json::to_string(&databricks_messages_with).unwrap(); let json_with = serde_json::to_string(&databricks_messages_with).unwrap();
println!("JSON with cache_control: {}", json_with); println!("JSON with cache_control: {}", json_with);
assert!(!json_with.contains("cache_control"), assert!(
"JSON should NOT contain 'cache_control' field - Databricks doesn't support it"); !json_with.contains("cache_control"),
"JSON should NOT contain 'cache_control' field - Databricks doesn't support it"
);
} }
#[test] #[test]
@@ -1343,7 +1355,13 @@ mod tests {
) )
.unwrap(); .unwrap();
assert!(!claude_provider.supports_cache_control(), "Databricks should not support cache_control even for Claude models"); assert!(
assert!(!llama_provider.supports_cache_control(), "Databricks should not support cache_control for Llama models"); !claude_provider.supports_cache_control(),
"Databricks should not support cache_control even for Claude models"
);
assert!(
!llama_provider.supports_cache_control(),
"Databricks should not support cache_control for Llama models"
);
} }
} }
+86 -40
View File
@@ -1,8 +1,8 @@
use anyhow::Result;
use crate::{ use crate::{
CompletionChunk, CompletionRequest, CompletionResponse, CompletionStream, LLMProvider, Message, CompletionChunk, CompletionRequest, CompletionResponse, CompletionStream, LLMProvider, Message,
MessageRole, Usage, MessageRole, Usage,
}; };
use anyhow::Result;
use llama_cpp::{ use llama_cpp::{
standard_sampler::{SamplerStage, StandardSampler}, standard_sampler::{SamplerStage, StandardSampler},
LlamaModel, LlamaParams, LlamaSession, SessionParams, LlamaModel, LlamaParams, LlamaSession, SessionParams,
@@ -148,7 +148,10 @@ impl EmbeddedProvider {
} }
// If the last message was from user, add a space for the assistant's response // If the last message was from user, add a space for the assistant's response
if messages.last().is_some_and(|m| matches!(m.role, MessageRole::User)) { if messages
.last()
.is_some_and(|m| matches!(m.role, MessageRole::User))
{
formatted.push(' '); formatted.push(' ');
} }
@@ -216,16 +219,25 @@ impl EmbeddedProvider {
} }
Err(_) => { Err(_) => {
if attempt < 4 { if attempt < 4 {
debug!("Session busy, retrying in {}ms (attempt {}/5)", 100 * (attempt + 1), attempt + 1); debug!(
std::thread::sleep(std::time::Duration::from_millis(100 * (attempt + 1) as u64)); "Session busy, retrying in {}ms (attempt {}/5)",
100 * (attempt + 1),
attempt + 1
);
std::thread::sleep(std::time::Duration::from_millis(
100 * (attempt + 1) as u64,
));
} else { } else {
return Err(anyhow::anyhow!("Model is busy after 5 attempts, please try again")); return Err(anyhow::anyhow!(
"Model is busy after 5 attempts, please try again"
));
} }
} }
} }
} }
let mut session = session_guard.ok_or_else(|| anyhow::anyhow!("Failed to acquire session lock"))?; let mut session = session_guard
.ok_or_else(|| anyhow::anyhow!("Failed to acquire session lock"))?;
debug!( debug!(
"Starting inference with prompt length: {} chars, estimated {} tokens", "Starting inference with prompt length: {} chars, estimated {} tokens",
@@ -350,18 +362,19 @@ impl EmbeddedProvider {
if model_name_lower.contains("qwen") { if model_name_lower.contains("qwen") {
vec![ vec![
"<|im_end|>", // Qwen ChatML format end token "<|im_end|>", // Qwen ChatML format end token
"<|endoftext|>", // Alternative end token "<|endoftext|>", // Alternative end token
"</s>", // Generic end of sequence "</s>", // Generic end of sequence
"<|im_start|>", // Start of new message (shouldn't appear in response) "<|im_start|>", // Start of new message (shouldn't appear in response)
] ]
} else if model_name_lower.contains("codellama") || model_name_lower.contains("code-llama") { } else if model_name_lower.contains("codellama") || model_name_lower.contains("code-llama")
{
vec![ vec![
"</s>", // End of sequence "</s>", // End of sequence
"[/INST]", // End of instruction "[/INST]", // End of instruction
"<</SYS>>", // End of system message "<</SYS>>", // End of system message
"[INST]", // Start of new instruction (shouldn't appear in response) "[INST]", // Start of new instruction (shouldn't appear in response)
"<<SYS>>", // Start of system (shouldn't appear in response) "<<SYS>>", // Start of system (shouldn't appear in response)
] ]
} else if model_name_lower.contains("llama") { } else if model_name_lower.contains("llama") {
vec![ vec![
@@ -374,9 +387,9 @@ impl EmbeddedProvider {
] ]
} else if model_name_lower.contains("mistral") { } else if model_name_lower.contains("mistral") {
vec![ vec![
"</s>", // End of sequence "</s>", // End of sequence
"[/INST]", // End of instruction "[/INST]", // End of instruction
"<|im_end|>", // ChatML format "<|im_end|>", // ChatML format
] ]
} else if model_name_lower.contains("vicuna") || model_name_lower.contains("wizard") { } else if model_name_lower.contains("vicuna") || model_name_lower.contains("wizard") {
vec![ vec![
@@ -391,7 +404,7 @@ impl EmbeddedProvider {
"### Instruction:", // Alpaca format "### Instruction:", // Alpaca format
"### Response:", // Alpaca format "### Response:", // Alpaca format
"### Input:", // Alpaca format "### Input:", // Alpaca format
"</s>", // End of sequence "</s>", // End of sequence
] ]
} else { } else {
// Generic/unknown model - use common stop sequences // Generic/unknown model - use common stop sequences
@@ -442,10 +455,11 @@ impl EmbeddedProvider {
// Use curl with progress bar for download // Use curl with progress bar for download
let output = Command::new("curl") let output = Command::new("curl")
.args([ .args([
"-L", // Follow redirects "-L", // Follow redirects
"-#", // Show progress bar "-#", // Show progress bar
"-f", // Fail on HTTP errors "-f", // Fail on HTTP errors
"-o", model_path.to_str().unwrap(), "-o",
model_path.to_str().unwrap(),
MODEL_URL, MODEL_URL,
]) ])
.output()?; .output()?;
@@ -455,11 +469,15 @@ impl EmbeddedProvider {
// If curl is not available, provide alternative instructions // If curl is not available, provide alternative instructions
if stderr.contains("command not found") || stderr.contains("not found") { if stderr.contains("command not found") || stderr.contains("not found") {
error!("curl is not installed. Please install curl or manually download the model."); error!(
"curl is not installed. Please install curl or manually download the model."
);
error!("Manual download instructions:"); error!("Manual download instructions:");
error!("1. Download from: {}", MODEL_URL); error!("1. Download from: {}", MODEL_URL);
error!("2. Save to: {}", model_path.display()); error!("2. Save to: {}", model_path.display());
anyhow::bail!("curl not found - please install curl or download the model manually"); anyhow::bail!(
"curl not found - please install curl or download the model manually"
);
} }
anyhow::bail!("Failed to download model: {}", stderr); anyhow::bail!("Failed to download model: {}", stderr);
@@ -469,11 +487,13 @@ impl EmbeddedProvider {
let metadata = fs::metadata(model_path)?; let metadata = fs::metadata(model_path)?;
let size_mb = metadata.len() / (1024 * 1024); let size_mb = metadata.len() / (1024 * 1024);
if size_mb < MODEL_SIZE_MB - 100 { // Allow some variance if size_mb < MODEL_SIZE_MB - 100 {
fs::remove_file(model_path).ok(); // Clean up partial download // Allow some variance
fs::remove_file(model_path).ok(); // Clean up partial download
anyhow::bail!( anyhow::bail!(
"Downloaded file appears incomplete ({}MB vs expected ~{}MB). Please try again.", "Downloaded file appears incomplete ({}MB vs expected ~{}MB). Please try again.",
size_mb, MODEL_SIZE_MB size_mb,
MODEL_SIZE_MB
); );
} }
@@ -541,10 +561,18 @@ impl LLMProvider for EmbeddedProvider {
} }
Err(_) => { Err(_) => {
if attempt < 4 { if attempt < 4 {
debug!("Session busy, retrying in {}ms (attempt {}/5)", 100 * (attempt + 1), attempt + 1); debug!(
std::thread::sleep(std::time::Duration::from_millis(100 * (attempt + 1) as u64)); "Session busy, retrying in {}ms (attempt {}/5)",
100 * (attempt + 1),
attempt + 1
);
std::thread::sleep(std::time::Duration::from_millis(
100 * (attempt + 1) as u64,
));
} else { } else {
let _ = tx.blocking_send(Err(anyhow::anyhow!("Model is busy after 5 attempts, please try again"))); let _ = tx.blocking_send(Err(anyhow::anyhow!(
"Model is busy after 5 attempts, please try again"
)));
return; return;
} }
} }
@@ -554,7 +582,8 @@ impl LLMProvider for EmbeddedProvider {
let mut session = match session_guard { let mut session = match session_guard {
Some(ctx) => ctx, Some(ctx) => ctx,
None => { None => {
let _ = tx.blocking_send(Err(anyhow::anyhow!("Failed to acquire session lock"))); let _ =
tx.blocking_send(Err(anyhow::anyhow!("Failed to acquire session lock")));
return; return;
} }
}; };
@@ -595,10 +624,26 @@ impl LLMProvider for EmbeddedProvider {
vec!["<|im_end|>", "<|endoftext|>", "</s>", "<|im_start|>"] vec!["<|im_end|>", "<|endoftext|>", "</s>", "<|im_start|>"]
} else if prompt.contains("[INST]") || prompt.contains("<<SYS>>") { } else if prompt.contains("[INST]") || prompt.contains("<<SYS>>") {
// Llama/CodeLlama format detected // Llama/CodeLlama format detected
vec!["</s>", "[/INST]", "<</SYS>>", "[INST]", "<<SYS>>", "### Human:", "### Assistant:"] vec![
"</s>",
"[/INST]",
"<</SYS>>",
"[INST]",
"<<SYS>>",
"### Human:",
"### Assistant:",
]
} else { } else {
// Generic format // Generic format
vec!["</s>", "<|endoftext|>", "<|im_end|>", "### Human:", "### Assistant:", "[/INST]", "<</SYS>>"] vec![
"</s>",
"<|endoftext|>",
"<|im_end|>",
"### Human:",
"### Assistant:",
"[/INST]",
"<</SYS>>",
]
}; };
// Stream tokens with proper limits // Stream tokens with proper limits
@@ -622,9 +667,9 @@ impl LLMProvider for EmbeddedProvider {
if hit_stop { if hit_stop {
// Before stopping, check if there might be an incomplete tool call // Before stopping, check if there might be an incomplete tool call
// Look for JSON tool call patterns that might be cut off by the stop sequence // Look for JSON tool call patterns that might be cut off by the stop sequence
let has_potential_tool_call = accumulated_text.contains(r#"{"tool":"#) || let has_potential_tool_call = accumulated_text.contains(r#"{"tool":"#)
accumulated_text.contains(r#"{"{""tool"":"#) || || accumulated_text.contains(r#"{"{""tool"":"#)
accumulated_text.contains(r#"{{""tool"":"#); || accumulated_text.contains(r#"{{""tool"":"#);
if has_potential_tool_call { if has_potential_tool_call {
// Check if the tool call appears to be complete (has closing brace after the stop sequence) // Check if the tool call appears to be complete (has closing brace after the stop sequence)
@@ -711,7 +756,8 @@ impl LLMProvider for EmbeddedProvider {
if might_be_stop { if might_be_stop {
// Hold back tokens, but only for a limited buffer size // Hold back tokens, but only for a limited buffer size
if unsent_tokens.len() > 20 { // Don't hold back more than 20 characters if unsent_tokens.len() > 20 {
// Don't hold back more than 20 characters
// Send the oldest part and keep only the recent part that might be a stop sequence // Send the oldest part and keep only the recent part that might be a stop sequence
let to_send = &unsent_tokens[..unsent_tokens.len() - 10]; let to_send = &unsent_tokens[..unsent_tokens.len() - 10];
if !to_send.is_empty() { if !to_send.is_empty() {
@@ -755,7 +801,7 @@ impl LLMProvider for EmbeddedProvider {
let final_chunk = CompletionChunk { let final_chunk = CompletionChunk {
content: String::new(), content: String::new(),
finished: true, finished: true,
usage: None, // Embedded models calculate usage differently usage: None, // Embedded models calculate usage differently
tool_calls: None, tool_calls: None,
}; };
let _ = tx.blocking_send(Ok(final_chunk)); let _ = tx.blocking_send(Ok(final_chunk));
+86 -30
View File
@@ -1,7 +1,7 @@
use serde::{Deserialize, Serialize};
use anyhow::Result; use anyhow::Result;
use std::collections::HashMap;
use rand::Rng; use rand::Rng;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
/// Trait for LLM providers /// Trait for LLM providers
#[async_trait::async_trait] #[async_trait::async_trait]
@@ -60,15 +60,24 @@ pub enum CacheType {
impl CacheControl { impl CacheControl {
pub fn ephemeral() -> Self { pub fn ephemeral() -> Self {
Self { cache_type: CacheType::Ephemeral, ttl: None } Self {
cache_type: CacheType::Ephemeral,
ttl: None,
}
} }
pub fn five_minute() -> Self { pub fn five_minute() -> Self {
Self { cache_type: CacheType::Ephemeral, ttl: Some("5m".to_string()) } Self {
cache_type: CacheType::Ephemeral,
ttl: Some("5m".to_string()),
}
} }
pub fn one_hour() -> Self { pub fn one_hour() -> Self {
Self { cache_type: CacheType::Ephemeral, ttl: Some("1h".to_string()) } Self {
cache_type: CacheType::Ephemeral,
ttl: Some("1h".to_string()),
}
} }
} }
@@ -76,6 +85,7 @@ impl CacheControl {
pub struct Message { pub struct Message {
pub role: MessageRole, pub role: MessageRole,
pub content: String, pub content: String,
#[serde(skip)]
pub id: String, pub id: String,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub cache_control: Option<CacheControl>, pub cache_control: Option<CacheControl>,
@@ -110,7 +120,7 @@ pub struct CompletionChunk {
pub content: String, pub content: String,
pub finished: bool, pub finished: bool,
pub tool_calls: Option<Vec<ToolCall>>, pub tool_calls: Option<Vec<ToolCall>>,
pub usage: Option<Usage>, // Add usage tracking for streaming pub usage: Option<Usage>, // Add usage tracking for streaming
} }
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -168,7 +178,11 @@ impl Message {
} }
/// Create a new message with cache control /// Create a new message with cache control
pub fn with_cache_control(role: MessageRole, content: String, cache_control: CacheControl) -> Self { pub fn with_cache_control(
role: MessageRole,
content: String,
cache_control: CacheControl,
) -> Self {
Self { Self {
role, role,
content, content,
@@ -182,7 +196,7 @@ impl Message {
role: MessageRole, role: MessageRole,
content: String, content: String,
cache_control: CacheControl, cache_control: CacheControl,
provider: &dyn LLMProvider provider: &dyn LLMProvider,
) -> Self { ) -> Self {
if !provider.supports_cache_control() { if !provider.supports_cache_control() {
tracing::warn!( tracing::warn!(
@@ -257,8 +271,10 @@ mod tests {
let json = serde_json::to_string(&msg).unwrap(); let json = serde_json::to_string(&msg).unwrap();
println!("Message JSON without cache_control: {}", json); println!("Message JSON without cache_control: {}", json);
assert!(!json.contains("cache_control"), assert!(
"JSON should not contain 'cache_control' field when not configured"); !json.contains("cache_control"),
"JSON should not contain 'cache_control' field when not configured"
);
} }
#[test] #[test]
@@ -271,14 +287,22 @@ mod tests {
let json = serde_json::to_string(&msg).unwrap(); let json = serde_json::to_string(&msg).unwrap();
println!("Message JSON with cache_control: {}", json); println!("Message JSON with cache_control: {}", json);
assert!(json.contains("cache_control"), assert!(
"JSON should contain 'cache_control' field when configured"); json.contains("cache_control"),
assert!(json.contains("ephemeral"), "JSON should contain 'cache_control' field when configured"
"JSON should contain 'ephemeral' value"); );
assert!(json.contains("\"type\":"), assert!(
"JSON should contain 'type' field in cache_control"); json.contains("ephemeral"),
assert!(!json.contains("null"), "JSON should contain 'ephemeral' value"
"JSON should not contain null values"); );
assert!(
json.contains("\"type\":"),
"JSON should contain 'type' field in cache_control"
);
assert!(
!json.contains("null"),
"JSON should not contain null values"
);
} }
#[test] #[test]
@@ -291,9 +315,18 @@ mod tests {
let json = serde_json::to_string(&msg).unwrap(); let json = serde_json::to_string(&msg).unwrap();
println!("Message JSON with 5-minute cache_control: {}", json); println!("Message JSON with 5-minute cache_control: {}", json);
assert!(json.contains("cache_control"), "JSON should contain 'cache_control' field"); assert!(
assert!(json.contains("ephemeral"), "JSON should contain 'ephemeral' type"); json.contains("cache_control"),
assert!(json.contains("\"ttl\":\"5m\""), "JSON should contain ttl field with 5m value"); "JSON should contain 'cache_control' field"
);
assert!(
json.contains("ephemeral"),
"JSON should contain 'ephemeral' type"
);
assert!(
json.contains("\"ttl\":\"5m\""),
"JSON should contain ttl field with 5m value"
);
} }
#[test] #[test]
@@ -306,9 +339,18 @@ mod tests {
let json = serde_json::to_string(&msg).unwrap(); let json = serde_json::to_string(&msg).unwrap();
println!("Message JSON with 1-hour cache_control: {}", json); println!("Message JSON with 1-hour cache_control: {}", json);
assert!(json.contains("cache_control"), "JSON should contain 'cache_control' field"); assert!(
assert!(json.contains("ephemeral"), "JSON should contain 'ephemeral' type"); json.contains("cache_control"),
assert!(json.contains("\"ttl\":\"1h\""), "JSON should contain ttl field with 1h value"); "JSON should contain 'cache_control' field"
);
assert!(
json.contains("ephemeral"),
"JSON should contain 'ephemeral' type"
);
assert!(
json.contains("\"ttl\":\"1h\""),
"JSON should contain ttl field with 1h value"
);
} }
#[test] #[test]
@@ -324,12 +366,17 @@ mod tests {
// Check timestamp part is 6 digits // Check timestamp part is 6 digits
assert_eq!(parts[0].len(), 6, "Timestamp should be 6 digits (HHMMSS)"); assert_eq!(parts[0].len(), 6, "Timestamp should be 6 digits (HHMMSS)");
assert!(parts[0].chars().all(|c| c.is_ascii_digit()), "Timestamp should be all digits"); assert!(
parts[0].chars().all(|c| c.is_ascii_digit()),
"Timestamp should be all digits"
);
// Check random part is 3 alpha characters // Check random part is 3 alpha characters
assert_eq!(parts[1].len(), 3, "Random part should be 3 characters"); assert_eq!(parts[1].len(), 3, "Random part should be 3 characters");
assert!(parts[1].chars().all(|c| c.is_ascii_alphabetic()), assert!(
"Random part should be all alphabetic characters"); parts[1].chars().all(|c| c.is_ascii_alphabetic()),
"Random part should be all alphabetic characters"
);
} }
#[test] #[test]
@@ -348,7 +395,10 @@ mod tests {
let json = serde_json::to_string(&msg).unwrap(); let json = serde_json::to_string(&msg).unwrap();
println!("Message JSON: {}", json); println!("Message JSON: {}", json);
assert!(!json.contains("\"id\""), "JSON should not contain 'id' field"); assert!(
!json.contains("\"id\""),
"JSON should not contain 'id' field"
);
} }
#[test] #[test]
@@ -359,7 +409,13 @@ mod tests {
CacheControl::ephemeral(), CacheControl::ephemeral(),
); );
assert!(!msg.id.is_empty(), "Message with cache control should have an ID"); assert!(
assert!(msg.id.contains('-'), "Message ID should contain hyphen separator"); !msg.id.is_empty(),
"Message with cache control should have an ID"
);
assert!(
msg.id.contains('-'),
"Message ID should contain hyphen separator"
);
} }
} }
+26 -10
View File
@@ -10,8 +10,8 @@ use tokio_stream::wrappers::ReceiverStream;
use tracing::{debug, error}; use tracing::{debug, error};
use crate::{ use crate::{
CompletionChunk, CompletionRequest, CompletionResponse, CompletionStream, LLMProvider, CompletionChunk, CompletionRequest, CompletionResponse, CompletionStream, LLMProvider, Message,
Message, MessageRole, Tool, ToolCall, Usage, MessageRole, Tool, ToolCall, Usage,
}; };
#[derive(Clone)] #[derive(Clone)]
@@ -138,7 +138,8 @@ impl OpenAIProvider {
debug!("Received stream completion marker"); debug!("Received stream completion marker");
// Send final chunk with accumulated content and tool calls // Send final chunk with accumulated content and tool calls
if !accumulated_content.is_empty() || !current_tool_calls.is_empty() { if !accumulated_content.is_empty() || !current_tool_calls.is_empty()
{
let tool_calls = if current_tool_calls.is_empty() { let tool_calls = if current_tool_calls.is_empty() {
None None
} else { } else {
@@ -188,8 +189,9 @@ impl OpenAIProvider {
if let Some(index) = delta_tool_call.index { if let Some(index) = delta_tool_call.index {
// Ensure we have enough tool calls in our vector // Ensure we have enough tool calls in our vector
while current_tool_calls.len() <= index { while current_tool_calls.len() <= index {
current_tool_calls current_tool_calls.push(
.push(OpenAIStreamingToolCall::default()); OpenAIStreamingToolCall::default(),
);
} }
let tool_call = &mut current_tool_calls[index]; let tool_call = &mut current_tool_calls[index];
@@ -198,11 +200,14 @@ impl OpenAIProvider {
tool_call.id = Some(id.clone()); tool_call.id = Some(id.clone());
} }
if let Some(function) = &delta_tool_call.function { if let Some(function) =
&delta_tool_call.function
{
if let Some(name) = &function.name { if let Some(name) = &function.name {
tool_call.name = Some(name.clone()); tool_call.name = Some(name.clone());
} }
if let Some(arguments) = &function.arguments { if let Some(arguments) = &function.arguments
{
tool_call.arguments.push_str(arguments); tool_call.arguments.push_str(arguments);
} }
} }
@@ -291,7 +296,11 @@ impl LLMProvider for OpenAIProvider {
.text() .text()
.await .await
.unwrap_or_else(|_| "Unknown error".to_string()); .unwrap_or_else(|_| "Unknown error".to_string());
return Err(anyhow::anyhow!("OpenAI API error {}: {}", status, error_text)); return Err(anyhow::anyhow!(
"OpenAI API error {}: {}",
status,
error_text
));
} }
let openai_response: OpenAIResponse = response.json().await?; let openai_response: OpenAIResponse = response.json().await?;
@@ -334,7 +343,10 @@ impl LLMProvider for OpenAIProvider {
request.temperature, request.temperature,
); );
debug!("Sending streaming request to OpenAI API: model={}", self.model); debug!(
"Sending streaming request to OpenAI API: model={}",
self.model
);
let response = self let response = self
.client .client
@@ -350,7 +362,11 @@ impl LLMProvider for OpenAIProvider {
.text() .text()
.await .await
.unwrap_or_else(|_| "Unknown error".to_string()); .unwrap_or_else(|_| "Unknown error".to_string());
return Err(anyhow::anyhow!("OpenAI API error {}: {}", status, error_text)); return Err(anyhow::anyhow!(
"OpenAI API error {}: {}",
status,
error_text
));
} }
let stream = response.bytes_stream(); let stream = response.bytes_stream();
@@ -24,14 +24,20 @@ fn test_no_wrong_serialization_format() {
println!("Ephemeral message JSON: {}", json); println!("Ephemeral message JSON: {}", json);
// Should NOT contain the wrong format // Should NOT contain the wrong format
assert!(!json.contains("system.0.cache_control"), assert!(
"JSON should not contain 'system.0.cache_control' path"); !json.contains("system.0.cache_control"),
assert!(!json.contains("cache_control.ephemeral"), "JSON should not contain 'system.0.cache_control' path"
"JSON should not contain 'cache_control.ephemeral' path"); );
assert!(
!json.contains("cache_control.ephemeral"),
"JSON should not contain 'cache_control.ephemeral' path"
);
// Should contain the correct format // Should contain the correct format
assert!(json.contains(r#""cache_control":{"type":"ephemeral"}"#), assert!(
"JSON should contain correct cache_control format"); json.contains(r#""cache_control":{"type":"ephemeral"}"#),
"JSON should contain correct cache_control format"
);
} }
#[test] #[test]
@@ -46,16 +52,24 @@ fn test_five_minute_no_wrong_format() {
println!("5-minute message JSON: {}", json); println!("5-minute message JSON: {}", json);
// Should NOT contain the wrong format // Should NOT contain the wrong format
assert!(!json.contains("system.0.cache_control"), assert!(
"JSON should not contain 'system.0.cache_control' path"); !json.contains("system.0.cache_control"),
assert!(!json.contains("cache_control.ephemeral.ttl"), "JSON should not contain 'system.0.cache_control' path"
"JSON should not contain 'cache_control.ephemeral.ttl' path"); );
assert!(
!json.contains("cache_control.ephemeral.ttl"),
"JSON should not contain 'cache_control.ephemeral.ttl' path"
);
// Should contain the correct format with ttl as a direct field // Should contain the correct format with ttl as a direct field
assert!(json.contains(r#""type":"ephemeral""#), assert!(
"JSON should contain type field"); json.contains(r#""type":"ephemeral""#),
assert!(json.contains(r#""ttl":"5m""#), "JSON should contain type field"
"JSON should contain ttl field with value 5m"); );
assert!(
json.contains(r#""ttl":"5m""#),
"JSON should contain ttl field with value 5m"
);
} }
#[test] #[test]
@@ -70,16 +84,24 @@ fn test_one_hour_no_wrong_format() {
println!("1-hour message JSON: {}", json); println!("1-hour message JSON: {}", json);
// Should NOT contain the wrong format // Should NOT contain the wrong format
assert!(!json.contains("system.0.cache_control"), assert!(
"JSON should not contain 'system.0.cache_control' path"); !json.contains("system.0.cache_control"),
assert!(!json.contains("cache_control.ephemeral.ttl"), "JSON should not contain 'system.0.cache_control' path"
"JSON should not contain 'cache_control.ephemeral.ttl' path"); );
assert!(
!json.contains("cache_control.ephemeral.ttl"),
"JSON should not contain 'cache_control.ephemeral.ttl' path"
);
// Should contain the correct format with ttl as a direct field // Should contain the correct format with ttl as a direct field
assert!(json.contains(r#""type":"ephemeral""#), assert!(
"JSON should contain type field"); json.contains(r#""type":"ephemeral""#),
assert!(json.contains(r#""ttl":"1h""#), "JSON should contain type field"
"JSON should contain ttl field with value 1h"); );
assert!(
json.contains(r#""ttl":"1h""#),
"JSON should contain ttl field with value 1h"
);
} }
#[test] #[test]
@@ -90,12 +112,19 @@ fn test_cache_control_structure_is_flat() {
let cache_control = CacheControl::five_minute(); let cache_control = CacheControl::five_minute();
let json_value = serde_json::to_value(&cache_control).unwrap(); let json_value = serde_json::to_value(&cache_control).unwrap();
println!("Cache control as JSON value: {}", serde_json::to_string_pretty(&json_value).unwrap()); println!(
"Cache control as JSON value: {}",
serde_json::to_string_pretty(&json_value).unwrap()
);
let obj = json_value.as_object().expect("Should be an object"); let obj = json_value.as_object().expect("Should be an object");
// Should have exactly 2 keys at the top level // Should have exactly 2 keys at the top level
assert_eq!(obj.len(), 2, "Cache control should have exactly 2 top-level fields"); assert_eq!(
obj.len(),
2,
"Cache control should have exactly 2 top-level fields"
);
// Both 'type' and 'ttl' should be at the same level // Both 'type' and 'ttl' should be at the same level
assert!(obj.contains_key("type"), "Should have 'type' field"); assert!(obj.contains_key("type"), "Should have 'type' field");
@@ -113,18 +142,28 @@ fn test_ephemeral_cache_control_structure() {
let cache_control = CacheControl::ephemeral(); let cache_control = CacheControl::ephemeral();
let json_value = serde_json::to_value(&cache_control).unwrap(); let json_value = serde_json::to_value(&cache_control).unwrap();
println!("Ephemeral cache control as JSON value: {}", serde_json::to_string_pretty(&json_value).unwrap()); println!(
"Ephemeral cache control as JSON value: {}",
serde_json::to_string_pretty(&json_value).unwrap()
);
let obj = json_value.as_object().expect("Should be an object"); let obj = json_value.as_object().expect("Should be an object");
// Should have exactly 1 key (only 'type', no 'ttl') // Should have exactly 1 key (only 'type', no 'ttl')
assert_eq!(obj.len(), 1, "Ephemeral cache control should have exactly 1 top-level field"); assert_eq!(
obj.len(),
1,
"Ephemeral cache control should have exactly 1 top-level field"
);
// Should have 'type' field // Should have 'type' field
assert!(obj.contains_key("type"), "Should have 'type' field"); assert!(obj.contains_key("type"), "Should have 'type' field");
// Should NOT have 'ttl' field // Should NOT have 'ttl' field
assert!(!obj.contains_key("ttl"), "Ephemeral should not have 'ttl' field"); assert!(
!obj.contains_key("ttl"),
"Ephemeral should not have 'ttl' field"
);
// 'type' should be a string with value "ephemeral" // 'type' should be a string with value "ephemeral"
assert_eq!(obj["type"].as_str().unwrap(), "ephemeral"); assert_eq!(obj["type"].as_str().unwrap(), "ephemeral");
@@ -11,11 +11,17 @@ fn test_ephemeral_cache_control_serialization() {
let cache_control = CacheControl::ephemeral(); let cache_control = CacheControl::ephemeral();
let json = serde_json::to_value(&cache_control).unwrap(); let json = serde_json::to_value(&cache_control).unwrap();
println!("Ephemeral cache_control JSON: {}", serde_json::to_string(&json).unwrap()); println!(
"Ephemeral cache_control JSON: {}",
serde_json::to_string(&json).unwrap()
);
assert_eq!(json, json!({ assert_eq!(
"type": "ephemeral" json,
})); json!({
"type": "ephemeral"
})
);
// Verify no ttl field is present // Verify no ttl field is present
assert!(!json.as_object().unwrap().contains_key("ttl")); assert!(!json.as_object().unwrap().contains_key("ttl"));
@@ -26,12 +32,18 @@ fn test_five_minute_cache_control_serialization() {
let cache_control = CacheControl::five_minute(); let cache_control = CacheControl::five_minute();
let json = serde_json::to_value(&cache_control).unwrap(); let json = serde_json::to_value(&cache_control).unwrap();
println!("5-minute cache_control JSON: {}", serde_json::to_string(&json).unwrap()); println!(
"5-minute cache_control JSON: {}",
serde_json::to_string(&json).unwrap()
);
assert_eq!(json, json!({ assert_eq!(
"type": "ephemeral", json,
"ttl": "5m" json!({
})); "type": "ephemeral",
"ttl": "5m"
})
);
} }
#[test] #[test]
@@ -39,12 +51,18 @@ fn test_one_hour_cache_control_serialization() {
let cache_control = CacheControl::one_hour(); let cache_control = CacheControl::one_hour();
let json = serde_json::to_value(&cache_control).unwrap(); let json = serde_json::to_value(&cache_control).unwrap();
println!("1-hour cache_control JSON: {}", serde_json::to_string(&json).unwrap()); println!(
"1-hour cache_control JSON: {}",
serde_json::to_string(&json).unwrap()
);
assert_eq!(json, json!({ assert_eq!(
"type": "ephemeral", json,
"ttl": "1h" json!({
})); "type": "ephemeral",
"ttl": "1h"
})
);
} }
#[test] #[test]
@@ -56,9 +74,14 @@ fn test_message_with_ephemeral_cache_control() {
); );
let json = serde_json::to_value(&msg).unwrap(); let json = serde_json::to_value(&msg).unwrap();
println!("Message with ephemeral cache_control: {}", serde_json::to_string(&json).unwrap()); println!(
"Message with ephemeral cache_control: {}",
serde_json::to_string(&json).unwrap()
);
let cache_control = json.get("cache_control").expect("cache_control field should exist"); let cache_control = json
.get("cache_control")
.expect("cache_control field should exist");
assert_eq!(cache_control.get("type").unwrap(), "ephemeral"); assert_eq!(cache_control.get("type").unwrap(), "ephemeral");
assert!(!cache_control.as_object().unwrap().contains_key("ttl")); assert!(!cache_control.as_object().unwrap().contains_key("ttl"));
} }
@@ -72,9 +95,14 @@ fn test_message_with_five_minute_cache_control() {
); );
let json = serde_json::to_value(&msg).unwrap(); let json = serde_json::to_value(&msg).unwrap();
println!("Message with 5-minute cache_control: {}", serde_json::to_string(&json).unwrap()); println!(
"Message with 5-minute cache_control: {}",
serde_json::to_string(&json).unwrap()
);
let cache_control = json.get("cache_control").expect("cache_control field should exist"); let cache_control = json
.get("cache_control")
.expect("cache_control field should exist");
assert_eq!(cache_control.get("type").unwrap(), "ephemeral"); assert_eq!(cache_control.get("type").unwrap(), "ephemeral");
assert_eq!(cache_control.get("ttl").unwrap(), "5m"); assert_eq!(cache_control.get("ttl").unwrap(), "5m");
} }
@@ -88,9 +116,14 @@ fn test_message_with_one_hour_cache_control() {
); );
let json = serde_json::to_value(&msg).unwrap(); let json = serde_json::to_value(&msg).unwrap();
println!("Message with 1-hour cache_control: {}", serde_json::to_string(&json).unwrap()); println!(
"Message with 1-hour cache_control: {}",
serde_json::to_string(&json).unwrap()
);
let cache_control = json.get("cache_control").expect("cache_control field should exist"); let cache_control = json
.get("cache_control")
.expect("cache_control field should exist");
assert_eq!(cache_control.get("type").unwrap(), "ephemeral"); assert_eq!(cache_control.get("type").unwrap(), "ephemeral");
assert_eq!(cache_control.get("ttl").unwrap(), "1h"); assert_eq!(cache_control.get("ttl").unwrap(), "1h");
} }
@@ -100,7 +133,10 @@ fn test_message_without_cache_control() {
let msg = Message::new(MessageRole::User, "Hello".to_string()); let msg = Message::new(MessageRole::User, "Hello".to_string());
let json = serde_json::to_value(&msg).unwrap(); let json = serde_json::to_value(&msg).unwrap();
println!("Message without cache_control: {}", serde_json::to_string(&json).unwrap()); println!(
"Message without cache_control: {}",
serde_json::to_string(&json).unwrap()
);
// cache_control field should not be present when not set // cache_control field should not be present when not set
assert!(!json.as_object().unwrap().contains_key("cache_control")); assert!(!json.as_object().unwrap().contains_key("cache_control"));
+23
View File
@@ -0,0 +1,23 @@
#!/bin/bash
# Hacky script for viewing context window
if [[ -n "$G3_WORKSPACE" ]]; then
TARGET_DIR="$G3_WORKSPACE/logs"
else
TARGET_DIR="$HOME/tmp/workspace/logs"
fi
if [[ ! -d "$TARGET_DIR" ]]; then
echo "Error: Directory '$TARGET_DIR' does not exist."
exit 1
fi
cd "$TARGET_DIR" || exit 1
NAME="$TARGET_DIR/current_context_window"
echo "Monitoring directory '$NAME' for current context window, (waits for first update)"
L=$(stat -f %m $NAME); while sleep 0.5; do N=$(stat -f %m $NAME); if [ "$N" != "$L" ]; then clear; cat $NAME; L=$N; fi; done