alex/g3
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add context window monitor

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

315
crates/g3-ensembles/src/flock.rs
View File

@@ -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");