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
@@ -17,19 +17,19 @@ use crate::status::{FlockStatus, SegmentState, SegmentStatus};
|
||||
pub struct FlockConfig {
|
||||
/// Project directory (must be a git repo with flock-requirements.md)
|
||||
pub project_dir: PathBuf,
|
||||
|
||||
|
||||
/// Flock workspace directory where segments will be created
|
||||
pub flock_workspace: PathBuf,
|
||||
|
||||
|
||||
/// Number of segments to partition work into
|
||||
pub num_segments: usize,
|
||||
|
||||
|
||||
/// Maximum turns per segment (for autonomous mode)
|
||||
pub max_turns: usize,
|
||||
|
||||
|
||||
/// G3 configuration to use for agents
|
||||
pub g3_config: Config,
|
||||
|
||||
|
||||
/// Path to g3 binary (defaults to current executable)
|
||||
pub g3_binary: Option<PathBuf>,
|
||||
}
|
||||
@@ -43,14 +43,20 @@ impl FlockConfig {
|
||||
) -> Result<Self> {
|
||||
// Validate project directory
|
||||
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
|
||||
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
|
||||
let requirements_path = project_dir.join("flock-requirements.md");
|
||||
if !requirements_path.exists() {
|
||||
@@ -59,10 +65,10 @@ impl FlockConfig {
|
||||
project_dir.display()
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
// Load default config
|
||||
let g3_config = Config::load(None)?;
|
||||
|
||||
|
||||
Ok(Self {
|
||||
project_dir,
|
||||
flock_workspace,
|
||||
@@ -72,19 +78,19 @@ impl FlockConfig {
|
||||
g3_binary: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
/// Set maximum turns per segment
|
||||
pub fn with_max_turns(mut self, max_turns: usize) -> Self {
|
||||
self.max_turns = max_turns;
|
||||
self
|
||||
}
|
||||
|
||||
|
||||
/// Set custom g3 binary path
|
||||
pub fn with_g3_binary(mut self, binary: PathBuf) -> Self {
|
||||
self.g3_binary = Some(binary);
|
||||
self
|
||||
}
|
||||
|
||||
|
||||
/// Set custom g3 config
|
||||
pub fn with_config(mut self, config: Config) -> Self {
|
||||
self.g3_config = config;
|
||||
@@ -103,58 +109,67 @@ impl FlockMode {
|
||||
/// Create a new flock mode instance
|
||||
pub fn new(config: FlockConfig) -> Result<Self> {
|
||||
let session_id = Uuid::new_v4().to_string();
|
||||
|
||||
|
||||
let status = FlockStatus::new(
|
||||
session_id.clone(),
|
||||
config.project_dir.clone(),
|
||||
config.flock_workspace.clone(),
|
||||
config.num_segments,
|
||||
);
|
||||
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
status,
|
||||
session_id,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
/// Run flock mode
|
||||
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
|
||||
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?;
|
||||
|
||||
|
||||
// Step 2: Create segment workspaces
|
||||
println!("\n📁 Step 2: Creating segment workspaces...");
|
||||
self.create_segment_workspaces(&partitions).await?;
|
||||
|
||||
|
||||
// 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?;
|
||||
|
||||
|
||||
// Step 4: Generate final report
|
||||
println!("\n📊 Step 4: Generating final report...");
|
||||
self.status.completed_at = Some(Utc::now());
|
||||
self.save_status()?;
|
||||
|
||||
|
||||
let report = self.status.generate_report();
|
||||
println!("{}", report);
|
||||
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
/// Partition requirements using an AI agent
|
||||
async fn partition_requirements(&mut self) -> Result<Vec<String>> {
|
||||
let requirements_path = self.config.project_dir.join("flock-requirements.md");
|
||||
let requirements_content = std::fs::read_to_string(&requirements_path)
|
||||
.context("Failed to read flock-requirements.md")?;
|
||||
|
||||
|
||||
// Create a temporary workspace for the partitioning agent
|
||||
let partition_workspace = self.config.flock_workspace.join("_partition");
|
||||
std::fs::create_dir_all(&partition_workspace)?;
|
||||
|
||||
|
||||
// Create the partitioning prompt
|
||||
let partition_prompt = format!(
|
||||
"You are a software architect tasked with partitioning project requirements into {} logical, \
|
||||
@@ -198,10 +213,10 @@ impl FlockMode {
|
||||
requirements_content,
|
||||
self.config.num_segments
|
||||
);
|
||||
|
||||
|
||||
// Get g3 binary path
|
||||
let g3_binary = self.get_g3_binary()?;
|
||||
|
||||
|
||||
// Run g3 in single-shot mode to partition requirements
|
||||
println!(" Analyzing requirements and creating partitions...");
|
||||
let output = Command::new(&g3_binary)
|
||||
@@ -212,23 +227,23 @@ impl FlockMode {
|
||||
.output()
|
||||
.await
|
||||
.context("Failed to run g3 for partitioning")?;
|
||||
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
anyhow::bail!("Partitioning agent failed: {}", stderr);
|
||||
}
|
||||
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
debug!("Partitioning agent output: {}", stdout);
|
||||
|
||||
|
||||
// Extract JSON from the output
|
||||
let partitions_json = Self::extract_json_from_output(&stdout)
|
||||
.context("Failed to extract partition JSON from agent output")?;
|
||||
|
||||
|
||||
// Parse the partitions
|
||||
let partitions: Vec<serde_json::Value> = serde_json::from_str(&partitions_json)
|
||||
.context("Failed to parse partition JSON")?;
|
||||
|
||||
let partitions: Vec<serde_json::Value> =
|
||||
serde_json::from_str(&partitions_json).context("Failed to parse partition JSON")?;
|
||||
|
||||
if partitions.len() != self.config.num_segments {
|
||||
warn!(
|
||||
"Expected {} partitions but got {}. Adjusting...",
|
||||
@@ -236,14 +251,12 @@ impl FlockMode {
|
||||
partitions.len()
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
// Extract requirements text from each partition
|
||||
let mut partition_texts = Vec::new();
|
||||
for (i, partition) in partitions.iter().enumerate() {
|
||||
let default_name = format!("module-{}", i + 1);
|
||||
let module_name = partition["module_name"]
|
||||
.as_str()
|
||||
.unwrap_or(&default_name);
|
||||
let module_name = partition["module_name"].as_str().unwrap_or(&default_name);
|
||||
let requirements = partition["requirements"]
|
||||
.as_str()
|
||||
.context("Missing requirements field in partition")?;
|
||||
@@ -256,7 +269,7 @@ impl FlockMode {
|
||||
.join(", ")
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
|
||||
let partition_text = format!(
|
||||
"# Module: {}\n\n## Dependencies\n{}\n\n## Requirements\n\n{}",
|
||||
module_name,
|
||||
@@ -267,69 +280,80 @@ impl FlockMode {
|
||||
},
|
||||
requirements
|
||||
);
|
||||
|
||||
|
||||
partition_texts.push(partition_text);
|
||||
println!(" ✓ Created partition {}: {}", i + 1, module_name);
|
||||
}
|
||||
|
||||
|
||||
Ok(partition_texts)
|
||||
}
|
||||
|
||||
|
||||
/// Extract JSON from agent output (looks for JSON array in output)
|
||||
fn extract_json_from_output(output: &str) -> Result<String> {
|
||||
// Try to find all occurrences of partition markers and extract valid JSON
|
||||
const MARKERS: &[&str] = &["{{PARTITION JSON}}", "{PARTITION JSON}"];
|
||||
|
||||
|
||||
let mut candidates = Vec::new();
|
||||
|
||||
|
||||
// Find all marker occurrences
|
||||
for &marker in MARKERS {
|
||||
let mut search_start = 0;
|
||||
while let Some(marker_index) = output[search_start..].find(marker) {
|
||||
let absolute_index = search_start + marker_index;
|
||||
let after_marker = &output[absolute_index + marker.len()..];
|
||||
|
||||
|
||||
// Try to find a code fence after this marker
|
||||
if let Some(fence_start) = after_marker.find("```") {
|
||||
let after_fence = &after_marker[fence_start + 3..];
|
||||
|
||||
|
||||
// Skip optional "json" language identifier
|
||||
let content_start = after_fence
|
||||
.strip_prefix("json")
|
||||
.unwrap_or(after_fence)
|
||||
.trim_start_matches(|c: char| c.is_whitespace());
|
||||
|
||||
|
||||
// Find closing fence
|
||||
if let Some(fence_end) = content_start.find("```") {
|
||||
let json_candidate = content_start[..fence_end].trim();
|
||||
candidates.push(json_candidate.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// Move search position forward
|
||||
search_start = absolute_index + marker.len();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
let mut last_error = None;
|
||||
for (i, candidate) in candidates.iter().enumerate() {
|
||||
match serde_json::from_str::<serde_json::Value>(candidate) {
|
||||
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());
|
||||
}
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// If we get here, none of the candidates were valid JSON
|
||||
if let Some(err) = last_error {
|
||||
anyhow::bail!(
|
||||
@@ -338,37 +362,46 @@ impl FlockMode {
|
||||
err
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
anyhow::bail!("No valid JSON found in output")
|
||||
}
|
||||
|
||||
|
||||
/// Create segment workspaces by copying project directory
|
||||
async fn create_segment_workspaces(&mut self, partitions: &[String]) -> Result<()> {
|
||||
// Ensure flock workspace exists
|
||||
std::fs::create_dir_all(&self.config.flock_workspace)?;
|
||||
|
||||
|
||||
for (i, partition) in partitions.iter().enumerate() {
|
||||
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);
|
||||
|
||||
|
||||
// Copy project directory to segment directory
|
||||
self.copy_git_repo(&self.config.project_dir, &segment_dir)
|
||||
.await
|
||||
.context(format!("Failed to copy project to segment {}", segment_id))?;
|
||||
|
||||
|
||||
// Write segment-requirements.md
|
||||
let requirements_path = segment_dir.join("segment-requirements.md");
|
||||
std::fs::write(&requirements_path, partition)
|
||||
.context(format!("Failed to write requirements for segment {}", segment_id))?;
|
||||
|
||||
println!(" ✓ Segment {} workspace ready at {}", segment_id, segment_dir.display());
|
||||
std::fs::write(&requirements_path, partition).context(format!(
|
||||
"Failed to write requirements for segment {}",
|
||||
segment_id
|
||||
))?;
|
||||
|
||||
println!(
|
||||
" ✓ Segment {} workspace ready at {}",
|
||||
segment_id,
|
||||
segment_dir.display()
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
/// Copy a git repository to a new location
|
||||
async fn copy_git_repo(&self, source: &Path, dest: &Path) -> Result<()> {
|
||||
// Use git clone for efficient copying
|
||||
@@ -379,26 +412,29 @@ impl FlockMode {
|
||||
.output()
|
||||
.await
|
||||
.context("Failed to run git clone")?;
|
||||
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
anyhow::bail!("Git clone failed: {}", stderr);
|
||||
}
|
||||
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
/// Run all segments in parallel
|
||||
async fn run_segments_parallel(&mut self) -> Result<()> {
|
||||
let mut handles = Vec::new();
|
||||
|
||||
|
||||
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 g3_binary = self.get_g3_binary()?;
|
||||
let status_file = self.get_status_file_path();
|
||||
let session_id = self.session_id.clone();
|
||||
|
||||
|
||||
// Initialize segment status
|
||||
let segment_status = SegmentStatus {
|
||||
segment_id,
|
||||
@@ -414,10 +450,10 @@ impl FlockMode {
|
||||
last_message: Some("Starting...".to_string()),
|
||||
error_message: None,
|
||||
};
|
||||
|
||||
|
||||
self.status.update_segment(segment_id, segment_status);
|
||||
self.save_status()?;
|
||||
|
||||
|
||||
// Spawn a task for this segment
|
||||
let handle = tokio::spawn(async move {
|
||||
run_segment(
|
||||
@@ -430,10 +466,10 @@ impl FlockMode {
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
|
||||
handles.push((segment_id, handle));
|
||||
}
|
||||
|
||||
|
||||
// Wait for all segments to complete
|
||||
for (segment_id, handle) in handles {
|
||||
match handle.await {
|
||||
@@ -444,10 +480,17 @@ impl FlockMode {
|
||||
}
|
||||
Ok(Err(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 {
|
||||
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,
|
||||
started_at: Utc::now(),
|
||||
completed_at: Some(Utc::now()),
|
||||
@@ -468,10 +511,17 @@ impl FlockMode {
|
||||
}
|
||||
Err(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 {
|
||||
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,
|
||||
started_at: Utc::now(),
|
||||
completed_at: Some(Utc::now()),
|
||||
@@ -492,10 +542,10 @@ impl FlockMode {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
/// Get the g3 binary path
|
||||
fn get_g3_binary(&self) -> Result<PathBuf> {
|
||||
if let Some(ref binary) = self.config.g3_binary {
|
||||
@@ -505,12 +555,12 @@ impl FlockMode {
|
||||
std::env::current_exe().context("Failed to get current executable path")
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Get the status file path
|
||||
fn get_status_file_path(&self) -> PathBuf {
|
||||
self.config.flock_workspace.join("flock-status.json")
|
||||
}
|
||||
|
||||
|
||||
/// Save current status to file
|
||||
fn save_status(&self) -> Result<()> {
|
||||
let status_file = self.get_status_file_path();
|
||||
@@ -527,8 +577,12 @@ async fn run_segment(
|
||||
status_file: PathBuf,
|
||||
session_id: String,
|
||||
) -> 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 {
|
||||
segment_id,
|
||||
workspace: segment_dir.clone(),
|
||||
@@ -543,7 +597,7 @@ async fn run_segment(
|
||||
last_message: Some("Starting autonomous mode...".to_string()),
|
||||
error_message: None,
|
||||
};
|
||||
|
||||
|
||||
// Run g3 in autonomous mode with segment-requirements.md
|
||||
let mut child = Command::new(&g3_binary)
|
||||
.arg("--workspace")
|
||||
@@ -552,23 +606,25 @@ async fn run_segment(
|
||||
.arg("--max-turns")
|
||||
.arg(max_turns.to_string())
|
||||
.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
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.context("Failed to spawn g3 process")?;
|
||||
|
||||
|
||||
// Stream output and update status
|
||||
let stdout = child.stdout.take().context("Failed to get stdout")?;
|
||||
let stderr = child.stderr.take().context("Failed to get stderr")?;
|
||||
|
||||
|
||||
let stdout_reader = BufReader::new(stdout);
|
||||
let stderr_reader = BufReader::new(stderr);
|
||||
|
||||
|
||||
let mut stdout_lines = stdout_reader.lines();
|
||||
let mut stderr_lines = stderr_reader.lines();
|
||||
|
||||
|
||||
// Read output and update status
|
||||
loop {
|
||||
tokio::select! {
|
||||
@@ -576,7 +632,7 @@ async fn run_segment(
|
||||
match line {
|
||||
Ok(Some(line)) => {
|
||||
println!("[Segment {}] {}", segment_id, line);
|
||||
|
||||
|
||||
// Parse output for status updates
|
||||
if line.contains("TURN") {
|
||||
// Extract turn number if possible
|
||||
@@ -586,7 +642,7 @@ async fn run_segment(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
segment_status.last_message = Some(line);
|
||||
update_status_file(&status_file, &session_id, segment_status.clone())?;
|
||||
}
|
||||
@@ -613,12 +669,15 @@ async fn run_segment(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// 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());
|
||||
|
||||
|
||||
if status.success() {
|
||||
segment_status.state = SegmentState::Completed;
|
||||
segment_status.last_message = Some("Completed successfully".to_string());
|
||||
@@ -627,7 +686,7 @@ async fn run_segment(
|
||||
segment_status.error_message = Some(format!("Process exited with status: {}", status));
|
||||
segment_status.errors += 1;
|
||||
}
|
||||
|
||||
|
||||
// Try to extract metrics from session log if available
|
||||
let log_dir = segment_dir.join("logs");
|
||||
if log_dir.exists() {
|
||||
@@ -636,7 +695,9 @@ async fn run_segment(
|
||||
let path = entry.path();
|
||||
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_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
|
||||
if let Some(context) = log_json.get("context_window") {
|
||||
if let Some(cumulative) = context.get("cumulative_tokens") {
|
||||
@@ -645,7 +706,7 @@ async fn run_segment(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// Count tool calls from conversation history
|
||||
if let Some(context) = log_json.get("context_window") {
|
||||
if let Some(history) = context.get("conversation_history") {
|
||||
@@ -653,8 +714,7 @@ async fn run_segment(
|
||||
let tool_call_count = messages
|
||||
.iter()
|
||||
.filter(|msg| {
|
||||
msg.get("role")
|
||||
.and_then(|r| r.as_str())
|
||||
msg.get("role").and_then(|r| r.as_str())
|
||||
== Some("tool")
|
||||
})
|
||||
.count();
|
||||
@@ -668,9 +728,9 @@ async fn run_segment(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
update_status_file(&status_file, &session_id, segment_status.clone())?;
|
||||
|
||||
|
||||
Ok(segment_status)
|
||||
}
|
||||
|
||||
@@ -685,24 +745,19 @@ fn update_status_file(
|
||||
FlockStatus::load_from_file(status_file)?
|
||||
} else {
|
||||
// This shouldn't happen, but handle it gracefully
|
||||
FlockStatus::new(
|
||||
session_id.to_string(),
|
||||
PathBuf::new(),
|
||||
PathBuf::new(),
|
||||
0,
|
||||
)
|
||||
FlockStatus::new(session_id.to_string(), PathBuf::new(), PathBuf::new(), 0)
|
||||
};
|
||||
|
||||
|
||||
flock_status.update_segment(segment_status.segment_id, segment_status);
|
||||
flock_status.save_to_file(status_file)?;
|
||||
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::FlockMode;
|
||||
|
||||
|
||||
#[test]
|
||||
fn extract_json_from_output_handles_partition_marker_and_fences() {
|
||||
const NOISY_PREFIX: &str = concat!(
|
||||
@@ -730,7 +785,7 @@ mod tests {
|
||||
"## Module Partitioning\n",
|
||||
"\n"
|
||||
);
|
||||
|
||||
|
||||
let expected_json = r#"[
|
||||
{
|
||||
"module_name": "message-protocol",
|
||||
@@ -743,18 +798,18 @@ mod tests {
|
||||
"dependencies": ["message-protocol"]
|
||||
}
|
||||
]"#;
|
||||
|
||||
|
||||
let mut output = String::from(NOISY_PREFIX);
|
||||
output.push_str("{{PARTITION JSON}}\n```json\n");
|
||||
output.push_str(expected_json);
|
||||
output.push_str("```");
|
||||
|
||||
|
||||
let extracted = FlockMode::extract_json_from_output(&output)
|
||||
.expect("should extract JSON between markers");
|
||||
|
||||
|
||||
assert_eq!(extracted, expected_json);
|
||||
}
|
||||
|
||||
|
||||
#[test]
|
||||
fn extract_json_from_output_handles_multiple_markers_and_invalid_json() {
|
||||
// This is the actual output from the LLM that was failing
|
||||
@@ -891,19 +946,19 @@ The requirements have been partitioned into two logical, largely non-overlapping
|
||||
4. **Maintainability**: Changes to logging/monitoring don't affect core message handling
|
||||
5. **Scalability**: Observability could be extracted to a separate service for distributed systems
|
||||
6. **Dependency Direction**: Clean one-way dependency (observability → message-protocol) prevents circular dependencies"#;
|
||||
|
||||
|
||||
let extracted = FlockMode::extract_json_from_output(output)
|
||||
.expect("should extract valid JSON from output with multiple markers");
|
||||
|
||||
|
||||
// Should be able to parse as JSON
|
||||
let parsed: serde_json::Value = serde_json::from_str(&extracted)
|
||||
.expect("extracted content should be valid JSON");
|
||||
|
||||
let parsed: serde_json::Value =
|
||||
serde_json::from_str(&extracted).expect("extracted content should be valid JSON");
|
||||
|
||||
// Verify it's an array with 2 elements
|
||||
assert!(parsed.is_array());
|
||||
let arr = parsed.as_array().unwrap();
|
||||
assert_eq!(arr.len(), 2);
|
||||
|
||||
|
||||
// Verify the structure
|
||||
assert_eq!(arr[0]["module_name"], "message-protocol");
|
||||
assert_eq!(arr[1]["module_name"], "observability");
|
||||
|
||||
@@ -10,37 +10,37 @@ use std::path::PathBuf;
|
||||
pub struct SegmentStatus {
|
||||
/// Segment number
|
||||
pub segment_id: usize,
|
||||
|
||||
|
||||
/// Segment workspace directory
|
||||
pub workspace: PathBuf,
|
||||
|
||||
|
||||
/// Current state of the segment
|
||||
pub state: SegmentState,
|
||||
|
||||
|
||||
/// Start time
|
||||
pub started_at: DateTime<Utc>,
|
||||
|
||||
|
||||
/// Completion time (if finished)
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
|
||||
|
||||
/// Total tokens used
|
||||
pub tokens_used: u64,
|
||||
|
||||
|
||||
/// Number of tool calls made
|
||||
pub tool_calls: u64,
|
||||
|
||||
|
||||
/// Number of errors encountered
|
||||
pub errors: u64,
|
||||
|
||||
|
||||
/// Current turn number (for autonomous mode)
|
||||
pub current_turn: usize,
|
||||
|
||||
|
||||
/// Maximum turns allowed
|
||||
pub max_turns: usize,
|
||||
|
||||
|
||||
/// Last status message
|
||||
pub last_message: Option<String>,
|
||||
|
||||
|
||||
/// Error message (if failed)
|
||||
pub error_message: Option<String>,
|
||||
}
|
||||
@@ -50,16 +50,16 @@ pub struct SegmentStatus {
|
||||
pub enum SegmentState {
|
||||
/// Waiting to start
|
||||
Pending,
|
||||
|
||||
|
||||
/// Currently running
|
||||
Running,
|
||||
|
||||
|
||||
/// Completed successfully
|
||||
Completed,
|
||||
|
||||
|
||||
/// Failed with error
|
||||
Failed,
|
||||
|
||||
|
||||
/// Cancelled by user
|
||||
Cancelled,
|
||||
}
|
||||
@@ -81,31 +81,31 @@ impl std::fmt::Display for SegmentState {
|
||||
pub struct FlockStatus {
|
||||
/// Flock session ID
|
||||
pub session_id: String,
|
||||
|
||||
|
||||
/// Project directory
|
||||
pub project_dir: PathBuf,
|
||||
|
||||
|
||||
/// Flock workspace directory
|
||||
pub flock_workspace: PathBuf,
|
||||
|
||||
|
||||
/// Number of segments
|
||||
pub num_segments: usize,
|
||||
|
||||
|
||||
/// Start time
|
||||
pub started_at: DateTime<Utc>,
|
||||
|
||||
|
||||
/// Completion time (if finished)
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
|
||||
|
||||
/// Status of each segment
|
||||
pub segments: HashMap<usize, SegmentStatus>,
|
||||
|
||||
|
||||
/// Total tokens used across all segments
|
||||
pub total_tokens: u64,
|
||||
|
||||
|
||||
/// Total tool calls across all segments
|
||||
pub total_tool_calls: u64,
|
||||
|
||||
|
||||
/// Total errors across all segments
|
||||
pub total_errors: u64,
|
||||
}
|
||||
@@ -131,20 +131,20 @@ impl FlockStatus {
|
||||
total_errors: 0,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Update segment status
|
||||
pub fn update_segment(&mut self, segment_id: usize, status: SegmentStatus) {
|
||||
self.segments.insert(segment_id, status);
|
||||
self.recalculate_totals();
|
||||
}
|
||||
|
||||
|
||||
/// Recalculate total metrics
|
||||
fn recalculate_totals(&mut self) {
|
||||
self.total_tokens = self.segments.values().map(|s| s.tokens_used).sum();
|
||||
self.total_tool_calls = self.segments.values().map(|s| s.tool_calls).sum();
|
||||
self.total_errors = self.segments.values().map(|s| s.errors).sum();
|
||||
}
|
||||
|
||||
|
||||
/// Check if all segments are complete
|
||||
pub fn is_complete(&self) -> bool {
|
||||
self.segments.len() == self.num_segments
|
||||
@@ -155,86 +155,116 @@ impl FlockStatus {
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
/// Get count of segments by state
|
||||
pub fn count_by_state(&self, state: SegmentState) -> usize {
|
||||
self.segments.values().filter(|s| s.state == state).count()
|
||||
}
|
||||
|
||||
|
||||
/// Save status to file
|
||||
pub fn save_to_file(&self, path: &PathBuf) -> anyhow::Result<()> {
|
||||
let json = serde_json::to_string_pretty(self)?;
|
||||
std::fs::write(path, json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
/// Load status from file
|
||||
pub fn load_from_file(path: &PathBuf) -> anyhow::Result<Self> {
|
||||
let json = std::fs::read_to_string(path)?;
|
||||
let status = serde_json::from_str(&json)?;
|
||||
Ok(status)
|
||||
}
|
||||
|
||||
|
||||
/// Generate a summary report
|
||||
pub fn generate_report(&self) -> String {
|
||||
let mut report = String::new();
|
||||
|
||||
|
||||
report.push_str(&format!("\n{}", "=".repeat(80)));
|
||||
report.push_str(&format!("\n📊 FLOCK MODE SESSION REPORT"));
|
||||
report.push_str(&format!("\n{}", "=".repeat(80)));
|
||||
|
||||
|
||||
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🗂️ Workspace: {}", self.flock_workspace.display()));
|
||||
report.push_str(&format!(
|
||||
"\n🗂️ Workspace: {}",
|
||||
self.flock_workspace.display()
|
||||
));
|
||||
report.push_str(&format!("\n🔢 Segments: {}", self.num_segments));
|
||||
|
||||
|
||||
let duration = if let Some(completed) = self.completed_at {
|
||||
completed.signed_duration_since(self.started_at)
|
||||
} else {
|
||||
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
|
||||
report.push_str(&format!("\n\n📈 Segment Status:"));
|
||||
report.push_str(&format!("\n • Completed: {}", self.count_by_state(SegmentState::Completed)));
|
||||
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)));
|
||||
|
||||
report.push_str(&format!(
|
||||
"\n • Completed: {}",
|
||||
self.count_by_state(SegmentState::Completed)
|
||||
));
|
||||
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
|
||||
report.push_str(&format!("\n\n📊 Aggregate Metrics:"));
|
||||
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));
|
||||
|
||||
|
||||
// Per-segment details
|
||||
report.push_str(&format!("\n\n🔍 Segment Details:"));
|
||||
let mut segments: Vec<_> = self.segments.iter().collect();
|
||||
segments.sort_by_key(|(id, _)| *id);
|
||||
|
||||
|
||||
for (id, segment) in segments {
|
||||
report.push_str(&format!("\n\n Segment {}:", id));
|
||||
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 Tool Calls: {}", segment.tool_calls));
|
||||
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 {
|
||||
report.push_str(&format!("\n Last Message: {}", msg));
|
||||
}
|
||||
|
||||
|
||||
if let Some(ref err) = segment.error_message {
|
||||
report.push_str(&format!("\n Error: {}", err));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
report.push_str(&format!("\n\n{}", "=".repeat(80)));
|
||||
|
||||
|
||||
report
|
||||
}
|
||||
}
|
||||
|
||||
@@ -283,8 +283,7 @@ mod tests {
|
||||
assert!(json.contains("Completed"));
|
||||
|
||||
// Deserialize back
|
||||
let deserialized: FlockStatus =
|
||||
serde_json::from_str(&json).expect("Failed to deserialize");
|
||||
let deserialized: FlockStatus = serde_json::from_str(&json).expect("Failed to deserialize");
|
||||
assert_eq!(deserialized.session_id, "test-session");
|
||||
assert_eq!(deserialized.segments.len(), 1);
|
||||
assert_eq!(deserialized.total_tokens, 1000);
|
||||
|
||||
Reference in New Issue
Block a user