refactor(g3-core): extract provider_registration and session modules
Extract two focused modules from the monolithic lib.rs (3372 lines): 1. provider_registration.rs (233 lines) - Consolidates duplicated provider registration patterns - Single determine_providers_to_register() function for mode-based selection - Unified register_providers() async function for all provider types - Includes unit tests for registration logic 2. session.rs (394 lines) - Session ID generation (generate_session_id) - Context window persistence (save_context_window, write_context_window_summary) - Error logging (log_error_to_session) - Utility functions (format_token_count, token_indicator) - Session restoration helper (restore_from_session_log) - Includes comprehensive unit tests Also fixes: - Removed redundant tool_executed assignment that triggered unused warning - Removed unused Message import in session.rs Results: - lib.rs reduced from 3372 to 2976 lines (-396 lines, -11.7%) - All tests pass, no warnings - Behavior preserved (pure mechanical extraction) Agent: fowler
This commit is contained in:
Dhanji R. Prasanna
394
crates/g3-core/src/session.rs
Normal file
394
crates/g3-core/src/session.rs
Normal file
@@ -0,0 +1,394 @@
|
||||
//! Session management utilities for the Agent.
|
||||
//!
|
||||
//! This module handles session ID generation, context window persistence,
|
||||
//! and session logging. It extracts the pure utility functions and I/O
|
||||
//! operations from the Agent, keeping the Agent as a thin orchestrator.
|
||||
|
||||
use crate::context_window::ContextWindow;
|
||||
use crate::paths::{ensure_session_dir, get_context_summary_file, get_g3_dir, get_logs_dir, get_session_file};
|
||||
use g3_providers::MessageRole;
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::path::PathBuf;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use tracing::{debug, error};
|
||||
|
||||
/// Format token count in compact form (e.g., 1K, 2M, 100b, 200K)
|
||||
/// Clamps to 4 chars right-aligned.
|
||||
pub fn format_token_count(tokens: u32) -> String {
|
||||
let mut raw = if tokens >= 1_000_000_000 {
|
||||
format!("{}b", tokens / 1_000_000_000)
|
||||
} else if tokens >= 1_000_000 {
|
||||
format!("{}M", tokens / 1_000_000)
|
||||
} else if tokens >= 1_000 {
|
||||
format!("{}K", tokens / 1_000)
|
||||
} else {
|
||||
"0K".to_string()
|
||||
};
|
||||
|
||||
if raw.len() > 4 {
|
||||
raw.truncate(4);
|
||||
}
|
||||
|
||||
format!("{:>4}", raw)
|
||||
}
|
||||
|
||||
/// Pick a single Unicode indicator for token magnitude (maps to color bands).
|
||||
pub fn token_indicator(tokens: u32) -> &'static str {
|
||||
if tokens <= 1_000 {
|
||||
"🟢"
|
||||
} else if tokens <= 5_000 {
|
||||
"🟡"
|
||||
} else if tokens <= 10_000 {
|
||||
"🟠"
|
||||
} else if tokens <= 20_000 {
|
||||
"🔴"
|
||||
} else {
|
||||
"🟣"
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate a session ID based on description and optional agent name.
|
||||
///
|
||||
/// For agent mode, uses agent name as prefix.
|
||||
/// For regular mode, uses first 5 words of description.
|
||||
/// Appends a hash for uniqueness.
|
||||
pub fn generate_session_id(description: &str, agent_name: Option<&str>) -> String {
|
||||
// For agent mode, use agent name as prefix for clarity
|
||||
// For regular mode, use first 5 words of description
|
||||
let prefix = if let Some(name) = agent_name {
|
||||
name.to_string()
|
||||
} else {
|
||||
description
|
||||
.chars()
|
||||
.filter(|c| c.is_alphanumeric() || c.is_whitespace() || *c == '-' || *c == '_')
|
||||
.collect::<String>()
|
||||
.split_whitespace()
|
||||
.take(5)
|
||||
.collect::<Vec<_>>()
|
||||
.join("_")
|
||||
.to_lowercase()
|
||||
};
|
||||
|
||||
// Create a hash for uniqueness (description + agent name + timestamp)
|
||||
let mut hasher = DefaultHasher::new();
|
||||
description.hash(&mut hasher);
|
||||
if let Some(name) = agent_name {
|
||||
name.hash(&mut hasher);
|
||||
}
|
||||
let timestamp = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos())
|
||||
.unwrap_or(0);
|
||||
timestamp.hash(&mut hasher);
|
||||
let hash = hasher.finish();
|
||||
|
||||
// Format: prefix_hash
|
||||
format!("{}_{:x}", prefix, hash)
|
||||
}
|
||||
|
||||
/// Save the context window to a session file.
|
||||
///
|
||||
/// If session_id is provided, saves to `.g3/sessions/<session_id>/session.json`.
|
||||
/// Otherwise, falls back to `logs/g3_context_<timestamp>.json`.
|
||||
pub fn save_context_window(
|
||||
session_id: Option<&str>,
|
||||
context_window: &ContextWindow,
|
||||
status: &str,
|
||||
) {
|
||||
let timestamp = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
// Determine filename based on session ID
|
||||
let filename = if let Some(id) = session_id {
|
||||
// Ensure session directory exists
|
||||
if let Err(e) = ensure_session_dir(id) {
|
||||
error!("Failed to create session directory: {}", e);
|
||||
return;
|
||||
}
|
||||
get_session_file(id)
|
||||
} else {
|
||||
// Fallback to old logs/ directory for sessions without ID
|
||||
let logs_dir = get_logs_dir();
|
||||
if let Err(e) = std::fs::create_dir_all(&logs_dir) {
|
||||
error!("Failed to create logs directory: {}", e);
|
||||
return;
|
||||
}
|
||||
logs_dir.join(format!("g3_context_{}.json", timestamp))
|
||||
};
|
||||
|
||||
let context_data = serde_json::json!({
|
||||
"session_id": session_id,
|
||||
"timestamp": timestamp,
|
||||
"status": status,
|
||||
"context_window": {
|
||||
"used_tokens": context_window.used_tokens,
|
||||
"total_tokens": context_window.total_tokens,
|
||||
"percentage_used": context_window.percentage_used(),
|
||||
"conversation_history": context_window.conversation_history
|
||||
}
|
||||
});
|
||||
|
||||
match serde_json::to_string_pretty(&context_data) {
|
||||
Ok(json_content) => {
|
||||
if let Err(e) = std::fs::write(&filename, &json_content) {
|
||||
error!("Failed to save context window to {:?}: {}", &filename, e);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Failed to serialize context window: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Write a human-readable context window summary to file.
|
||||
///
|
||||
/// Format: message_id, role, token_count, indicator, first_120_chars
|
||||
pub fn write_context_window_summary(session_id: &str, context_window: &ContextWindow) {
|
||||
// Ensure session directory exists
|
||||
if let Err(e) = ensure_session_dir(session_id) {
|
||||
error!("Failed to create session directory: {}", e);
|
||||
return;
|
||||
}
|
||||
|
||||
let filename = get_context_summary_file(session_id);
|
||||
let symlink_path = get_g3_dir().join("sessions").join("current_context_window");
|
||||
|
||||
// Build the summary content
|
||||
let mut summary_lines = Vec::new();
|
||||
|
||||
for message in &context_window.conversation_history {
|
||||
// Estimate tokens for this message
|
||||
let message_tokens = ContextWindow::estimate_tokens(&message.content);
|
||||
|
||||
// Format token count and get indicator
|
||||
let token_str = format_token_count(message_tokens);
|
||||
let indicator = token_indicator(message_tokens);
|
||||
|
||||
// Get role as string
|
||||
let role = match message.role {
|
||||
MessageRole::System => "sys",
|
||||
MessageRole::User => "usr",
|
||||
MessageRole::Assistant => "ass",
|
||||
};
|
||||
|
||||
// Get first 120 characters of content, replace newlines
|
||||
let content_preview: String = message
|
||||
.content
|
||||
.chars()
|
||||
.take(120)
|
||||
.collect::<String>()
|
||||
.replace('\n', " ")
|
||||
.replace('\r', " ");
|
||||
|
||||
let line = format!(
|
||||
"{}, {}, {} {}, {}\n",
|
||||
message.id, role, token_str, indicator, content_preview
|
||||
);
|
||||
summary_lines.push(line);
|
||||
}
|
||||
|
||||
// Add total estimate
|
||||
let total_token_str = format_token_count(context_window.used_tokens);
|
||||
let capacity_str = format_token_count(context_window.total_tokens);
|
||||
let percentage = context_window.percentage_used();
|
||||
|
||||
summary_lines.push(format!(
|
||||
"\n--- TOTAL: {} / {} ({:.1}%) ---\n",
|
||||
total_token_str, capacity_str, percentage
|
||||
));
|
||||
|
||||
// Write to file
|
||||
let summary_content = summary_lines.join("");
|
||||
if let Err(e) = std::fs::write(&filename, &summary_content) {
|
||||
error!(
|
||||
"Failed to write context window summary to {:?}: {}",
|
||||
&filename, e
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// Update symlink
|
||||
let _ = std::fs::remove_file(&symlink_path);
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::symlink;
|
||||
let target = format!("context_window_{}.txt", session_id);
|
||||
if let Err(e) = symlink(&target, &symlink_path) {
|
||||
error!("Failed to create symlink {:?}: {}", &symlink_path, e);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
{
|
||||
use std::os::windows::fs::symlink_file;
|
||||
let target = format!("context_window_{}.txt", session_id);
|
||||
if let Err(e) = symlink_file(&target, &symlink_path) {
|
||||
error!("Failed to create symlink {:?}: {}", &symlink_path, e);
|
||||
}
|
||||
}
|
||||
|
||||
debug!(
|
||||
"Context window summary written to {:?} ({} messages)",
|
||||
filename,
|
||||
context_window.conversation_history.len()
|
||||
);
|
||||
}
|
||||
|
||||
/// Log an error to the session JSON file.
|
||||
///
|
||||
/// Appends an error entry to the conversation history in the session log.
|
||||
pub fn log_error_to_session(
|
||||
session_id: &str,
|
||||
error: &anyhow::Error,
|
||||
role: &str,
|
||||
forensic_context: Option<String>,
|
||||
) {
|
||||
let timestamp = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
let logs_dir = get_logs_dir();
|
||||
let filename = logs_dir.join(format!("g3_session_{}.json", session_id));
|
||||
|
||||
// Read existing session log
|
||||
let mut session_data: serde_json::Value = if filename.exists() {
|
||||
match std::fs::read_to_string(&filename) {
|
||||
Ok(content) => serde_json::from_str(&content).unwrap_or_else(|_| serde_json::json!({})),
|
||||
Err(_) => serde_json::json!({}),
|
||||
}
|
||||
} else {
|
||||
serde_json::json!({})
|
||||
};
|
||||
|
||||
// Build error message with forensic context
|
||||
let error_message = if let Some(context) = forensic_context {
|
||||
format!("ERROR: {}\n\nForensic Context:\n{}", error, context)
|
||||
} else {
|
||||
format!("ERROR: {}", error)
|
||||
};
|
||||
|
||||
// Create error message entry
|
||||
let error_entry = serde_json::json!({
|
||||
"role": role,
|
||||
"content": error_message,
|
||||
"timestamp": timestamp,
|
||||
"error_type": "context_length_exceeded"
|
||||
});
|
||||
|
||||
// Append to conversation history
|
||||
if let Some(history) = session_data
|
||||
.get_mut("context_window")
|
||||
.and_then(|cw| cw.get_mut("conversation_history"))
|
||||
{
|
||||
if let Some(history_array) = history.as_array_mut() {
|
||||
history_array.push(error_entry);
|
||||
}
|
||||
}
|
||||
|
||||
// Write back to file
|
||||
if let Ok(json_content) = serde_json::to_string_pretty(&session_data) {
|
||||
let _ = std::fs::write(&filename, json_content);
|
||||
}
|
||||
}
|
||||
|
||||
/// Restore conversation history from a session log file.
|
||||
///
|
||||
/// Returns the messages to add to the context window, or None if restoration failed.
|
||||
pub fn restore_from_session_log(session_log_path: &PathBuf) -> Option<Vec<(MessageRole, String)>> {
|
||||
if !session_log_path.exists() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let json = std::fs::read_to_string(session_log_path).ok()?;
|
||||
let session_data: serde_json::Value = serde_json::from_str(&json).ok()?;
|
||||
|
||||
let context_window = session_data.get("context_window")?;
|
||||
let history = context_window.get("conversation_history")?;
|
||||
let messages = history.as_array()?;
|
||||
|
||||
let mut result = Vec::new();
|
||||
for msg in messages {
|
||||
let role_str = msg.get("role").and_then(|r| r.as_str()).unwrap_or("user");
|
||||
let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
|
||||
|
||||
// Skip system messages (they're preserved separately)
|
||||
if role_str == "system" {
|
||||
continue;
|
||||
}
|
||||
|
||||
let role = match role_str {
|
||||
"assistant" => MessageRole::Assistant,
|
||||
_ => MessageRole::User,
|
||||
};
|
||||
|
||||
result.push((role, content.to_string()));
|
||||
}
|
||||
|
||||
Some(result)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_format_token_count_small() {
|
||||
assert_eq!(format_token_count(0), " 0K");
|
||||
assert_eq!(format_token_count(500), " 0K");
|
||||
assert_eq!(format_token_count(999), " 0K");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_token_count_thousands() {
|
||||
assert_eq!(format_token_count(1000), " 1K");
|
||||
assert_eq!(format_token_count(5000), " 5K");
|
||||
assert_eq!(format_token_count(10000), " 10K");
|
||||
assert_eq!(format_token_count(999999), "999K");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_token_count_millions() {
|
||||
assert_eq!(format_token_count(1_000_000), " 1M");
|
||||
assert_eq!(format_token_count(5_000_000), " 5M");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_indicator() {
|
||||
assert_eq!(token_indicator(500), "🟢");
|
||||
assert_eq!(token_indicator(1000), "🟢");
|
||||
assert_eq!(token_indicator(1001), "🟡");
|
||||
assert_eq!(token_indicator(5000), "🟡");
|
||||
assert_eq!(token_indicator(5001), "🟠");
|
||||
assert_eq!(token_indicator(10000), "🟠");
|
||||
assert_eq!(token_indicator(10001), "🔴");
|
||||
assert_eq!(token_indicator(20000), "🔴");
|
||||
assert_eq!(token_indicator(20001), "🟣");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_generate_session_id_regular_mode() {
|
||||
let id = generate_session_id("implement a function to calculate fibonacci", None);
|
||||
assert!(id.starts_with("implement_a_function_to_calculate_"));
|
||||
assert!(id.contains('_')); // Has hash suffix
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_generate_session_id_agent_mode() {
|
||||
let id = generate_session_id("some task", Some("fowler"));
|
||||
assert!(id.starts_with("fowler_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_generate_session_id_uniqueness() {
|
||||
// Same description should produce different IDs due to timestamp
|
||||
let id1 = generate_session_id("test", None);
|
||||
std::thread::sleep(std::time::Duration::from_millis(1));
|
||||
let id2 = generate_session_id("test", None);
|
||||
assert_ne!(id1, id2);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user