mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Add heartbeat integration, planning phase, and auto-repair
- Add HeartbeatConfig for proactive periodic execution with channel notifications - Add use_planning option to Worker for ActionPlan generation before tool execution - Implement tool failure tracking in database (V3 migration) - Add auto-repair via Builder for broken WASM tools in self_repair.rs - Record tool failures in Worker for self-repair tracking - Update .env.example with new configuration options Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
Illia Polosukhin
co-authored by
Claude Opus 4.5
Claude Opus 4.5
parent
2df4a4f5f0
commit
235f6aae18
+63
-2
@@ -8,13 +8,16 @@ use uuid::Uuid;
|
||||
|
||||
use crate::agent::compaction::ContextCompactor;
|
||||
use crate::agent::context_monitor::ContextMonitor;
|
||||
use crate::agent::heartbeat::spawn_heartbeat;
|
||||
use crate::agent::self_repair::DefaultSelfRepair;
|
||||
use crate::agent::session::{Session, ThreadState};
|
||||
use crate::agent::session_manager::SessionManager;
|
||||
use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult};
|
||||
use crate::agent::{MessageIntent, RepairTask, Router, Scheduler};
|
||||
use crate::agent::{
|
||||
HeartbeatConfig as AgentHeartbeatConfig, MessageIntent, RepairTask, Router, Scheduler,
|
||||
};
|
||||
use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusUpdate};
|
||||
use crate::config::AgentConfig;
|
||||
use crate::config::{AgentConfig, HeartbeatConfig};
|
||||
use crate::context::ContextManager;
|
||||
use crate::error::Error;
|
||||
use crate::history::Store;
|
||||
@@ -37,6 +40,7 @@ pub struct Agent {
|
||||
session_manager: Arc<SessionManager>,
|
||||
context_monitor: ContextMonitor,
|
||||
workspace: Option<Arc<Workspace>>,
|
||||
heartbeat_config: Option<HeartbeatConfig>,
|
||||
}
|
||||
|
||||
impl Agent {
|
||||
@@ -49,6 +53,7 @@ impl Agent {
|
||||
tools: Arc<ToolRegistry>,
|
||||
channels: ChannelManager,
|
||||
workspace: Option<Arc<Workspace>>,
|
||||
heartbeat_config: Option<HeartbeatConfig>,
|
||||
) -> Self {
|
||||
let context_manager = Arc::new(ContextManager::new(config.max_parallel_jobs));
|
||||
|
||||
@@ -74,6 +79,7 @@ impl Agent {
|
||||
session_manager: Arc::new(SessionManager::new()),
|
||||
context_monitor: ContextMonitor::new(),
|
||||
workspace,
|
||||
heartbeat_config,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -94,6 +100,58 @@ impl Agent {
|
||||
repair_task.run().await;
|
||||
});
|
||||
|
||||
// Spawn heartbeat if enabled
|
||||
let heartbeat_handle = if let Some(ref hb_config) = self.heartbeat_config {
|
||||
if hb_config.enabled {
|
||||
if let Some(ref workspace) = self.workspace {
|
||||
let config = AgentHeartbeatConfig::default()
|
||||
.with_interval(std::time::Duration::from_secs(hb_config.interval_secs));
|
||||
|
||||
// Set up notification channel if configured
|
||||
let (notify_tx, mut notify_rx) =
|
||||
tokio::sync::mpsc::channel::<OutgoingResponse>(16);
|
||||
|
||||
// Spawn notification forwarder
|
||||
// We can't clone ChannelManager directly, so we just log the notifications
|
||||
// The heartbeat system will handle notifications via the response_tx
|
||||
let notify_channel = hb_config.notify_channel.clone();
|
||||
let notify_user = hb_config.notify_user.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(response) = notify_rx.recv().await {
|
||||
if let (Some(ch), Some(user)) = (¬ify_channel, ¬ify_user) {
|
||||
// Log the heartbeat notification
|
||||
// In a full implementation, we'd route this through a shared channel reference
|
||||
tracing::info!(
|
||||
"Heartbeat notification for {}/{}: {}",
|
||||
ch,
|
||||
user,
|
||||
&response.content
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tracing::info!(
|
||||
"Heartbeat enabled with {}s interval",
|
||||
hb_config.interval_secs
|
||||
);
|
||||
Some(spawn_heartbeat(
|
||||
config,
|
||||
workspace.clone(),
|
||||
self.llm.clone(),
|
||||
Some(notify_tx),
|
||||
))
|
||||
} else {
|
||||
tracing::warn!("Heartbeat enabled but no workspace available");
|
||||
None
|
||||
}
|
||||
} else {
|
||||
None
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Main message loop
|
||||
tracing::info!("Agent {} ready and listening", self.config.name);
|
||||
|
||||
@@ -123,6 +181,9 @@ impl Agent {
|
||||
// Cleanup
|
||||
tracing::info!("Agent shutting down...");
|
||||
repair_handle.abort();
|
||||
if let Some(handle) = heartbeat_handle {
|
||||
handle.abort();
|
||||
}
|
||||
self.scheduler.stop_all().await;
|
||||
self.channels.shutdown_all().await?;
|
||||
|
||||
|
||||
+1
-1
@@ -29,7 +29,7 @@ pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor};
|
||||
pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat};
|
||||
pub use router::{MessageIntent, Router};
|
||||
pub use scheduler::Scheduler;
|
||||
pub use self_repair::{RepairResult, RepairTask, SelfRepair, StuckJob};
|
||||
pub use self_repair::{BrokenTool, RepairResult, RepairTask, SelfRepair, StuckJob};
|
||||
pub use session::{Session, Thread, ThreadState, Turn, TurnState};
|
||||
pub use session_manager::SessionManager;
|
||||
pub use submission::{Submission, SubmissionParser, SubmissionResult};
|
||||
|
||||
@@ -120,6 +120,7 @@ impl Scheduler {
|
||||
self.tools.clone(),
|
||||
self.store.clone(),
|
||||
self.config.job_timeout,
|
||||
self.config.use_planning,
|
||||
);
|
||||
|
||||
// Spawn worker task
|
||||
|
||||
+147
-9
@@ -9,6 +9,8 @@ use uuid::Uuid;
|
||||
|
||||
use crate::context::{ContextManager, JobState};
|
||||
use crate::error::RepairError;
|
||||
use crate::history::Store;
|
||||
use crate::tools::{BuildRequirement, Language, SoftwareBuilder, SoftwareType, ToolRegistry};
|
||||
|
||||
/// A job that has been detected as stuck.
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -26,7 +28,10 @@ pub struct BrokenTool {
|
||||
pub name: String,
|
||||
pub failure_count: u32,
|
||||
pub last_error: Option<String>,
|
||||
pub first_failure: DateTime<Utc>,
|
||||
pub last_failure: DateTime<Utc>,
|
||||
pub last_build_result: Option<serde_json::Value>,
|
||||
pub repair_attempts: u32,
|
||||
}
|
||||
|
||||
/// Result of a repair attempt.
|
||||
@@ -63,6 +68,9 @@ pub struct DefaultSelfRepair {
|
||||
context_manager: Arc<ContextManager>,
|
||||
stuck_threshold: Duration,
|
||||
max_repair_attempts: u32,
|
||||
store: Option<Arc<Store>>,
|
||||
builder: Option<Arc<dyn SoftwareBuilder>>,
|
||||
tools: Option<Arc<ToolRegistry>>,
|
||||
}
|
||||
|
||||
impl DefaultSelfRepair {
|
||||
@@ -76,8 +84,28 @@ impl DefaultSelfRepair {
|
||||
context_manager,
|
||||
stuck_threshold,
|
||||
max_repair_attempts,
|
||||
store: None,
|
||||
builder: None,
|
||||
tools: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a Store for tool failure tracking.
|
||||
pub fn with_store(mut self, store: Arc<Store>) -> Self {
|
||||
self.store = Some(store);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add a Builder and ToolRegistry for automatic tool repair.
|
||||
pub fn with_builder(
|
||||
mut self,
|
||||
builder: Arc<dyn SoftwareBuilder>,
|
||||
tools: Arc<ToolRegistry>,
|
||||
) -> Self {
|
||||
self.builder = Some(builder);
|
||||
self.tools = Some(tools);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -151,19 +179,129 @@ impl SelfRepair for DefaultSelfRepair {
|
||||
}
|
||||
|
||||
async fn detect_broken_tools(&self) -> Vec<BrokenTool> {
|
||||
// TODO: Implement tool failure tracking
|
||||
// Would need to track tool failures in the database
|
||||
vec![]
|
||||
let Some(ref store) = self.store else {
|
||||
return vec![];
|
||||
};
|
||||
|
||||
// Threshold: 5 failures before considering a tool broken
|
||||
match store.get_broken_tools(5).await {
|
||||
Ok(tools) => {
|
||||
if !tools.is_empty() {
|
||||
tracing::info!("Detected {} broken tools needing repair", tools.len());
|
||||
}
|
||||
tools
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to detect broken tools: {}", e);
|
||||
vec![]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn repair_broken_tool(&self, tool: &BrokenTool) -> Result<RepairResult, RepairError> {
|
||||
// TODO: Implement tool repair via ToolBuilder
|
||||
Ok(RepairResult::ManualRequired {
|
||||
message: format!(
|
||||
"Tool '{}' repair not implemented - manual intervention required",
|
||||
tool.name
|
||||
let Some(ref builder) = self.builder else {
|
||||
return Ok(RepairResult::ManualRequired {
|
||||
message: format!("Builder not available for repairing tool '{}'", tool.name),
|
||||
});
|
||||
};
|
||||
|
||||
let Some(ref store) = self.store else {
|
||||
return Ok(RepairResult::ManualRequired {
|
||||
message: "Store not available for tracking repair".to_string(),
|
||||
});
|
||||
};
|
||||
|
||||
// Check repair attempt limit
|
||||
if tool.repair_attempts >= self.max_repair_attempts {
|
||||
return Ok(RepairResult::ManualRequired {
|
||||
message: format!(
|
||||
"Tool '{}' exceeded max repair attempts ({})",
|
||||
tool.name, self.max_repair_attempts
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
"Attempting to repair tool '{}' (attempt {})",
|
||||
tool.name,
|
||||
tool.repair_attempts + 1
|
||||
);
|
||||
|
||||
// Increment repair attempts
|
||||
if let Err(e) = store.increment_repair_attempts(&tool.name).await {
|
||||
tracing::warn!("Failed to increment repair attempts: {}", e);
|
||||
}
|
||||
|
||||
// Create BuildRequirement for repair
|
||||
let requirement = BuildRequirement {
|
||||
name: tool.name.clone(),
|
||||
description: format!(
|
||||
"Repair broken WASM tool.\n\n\
|
||||
Tool name: {}\n\
|
||||
Previous error: {}\n\
|
||||
Failure count: {}\n\n\
|
||||
Analyze the error, fix the implementation, and rebuild.",
|
||||
tool.name,
|
||||
tool.last_error.as_deref().unwrap_or("Unknown error"),
|
||||
tool.failure_count
|
||||
),
|
||||
})
|
||||
software_type: SoftwareType::WasmTool,
|
||||
language: Language::Rust,
|
||||
input_spec: None,
|
||||
output_spec: None,
|
||||
dependencies: vec![],
|
||||
capabilities: vec!["http".to_string(), "workspace".to_string()],
|
||||
};
|
||||
|
||||
// Attempt to build/repair
|
||||
match builder.build(&requirement).await {
|
||||
Ok(result) if result.success => {
|
||||
tracing::info!(
|
||||
"Successfully rebuilt tool '{}' after {} iterations",
|
||||
tool.name,
|
||||
result.iterations
|
||||
);
|
||||
|
||||
// Mark as repaired in database
|
||||
if let Err(e) = store.mark_tool_repaired(&tool.name).await {
|
||||
tracing::warn!("Failed to mark tool as repaired: {}", e);
|
||||
}
|
||||
|
||||
// Log if the tool was auto-registered
|
||||
if result.registered {
|
||||
tracing::info!("Repaired tool '{}' auto-registered", tool.name);
|
||||
}
|
||||
|
||||
Ok(RepairResult::Success {
|
||||
message: format!(
|
||||
"Tool '{}' repaired successfully after {} iterations",
|
||||
tool.name, result.iterations
|
||||
),
|
||||
})
|
||||
}
|
||||
Ok(result) => {
|
||||
// Build completed but failed
|
||||
tracing::warn!(
|
||||
"Repair build for '{}' completed but failed: {:?}",
|
||||
tool.name,
|
||||
result.error
|
||||
);
|
||||
Ok(RepairResult::Retry {
|
||||
message: format!(
|
||||
"Repair attempt {} for '{}' failed: {}",
|
||||
tool.repair_attempts + 1,
|
||||
tool.name,
|
||||
result.error.unwrap_or_else(|| "Unknown error".to_string())
|
||||
),
|
||||
})
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Repair build for '{}' errored: {}", tool.name, e);
|
||||
Ok(RepairResult::Retry {
|
||||
message: format!("Repair build error: {}", e),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+153
-1
@@ -12,7 +12,9 @@ use crate::agent::task::TaskOutput;
|
||||
use crate::context::{ContextManager, JobState};
|
||||
use crate::error::Error;
|
||||
use crate::history::Store;
|
||||
use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, ToolSelection};
|
||||
use crate::llm::{
|
||||
ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, ToolSelection,
|
||||
};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::ToolRegistry;
|
||||
|
||||
@@ -25,6 +27,8 @@ pub struct Worker {
|
||||
tools: Arc<ToolRegistry>,
|
||||
store: Option<Arc<Store>>,
|
||||
timeout: Duration,
|
||||
/// Whether to use planning before tool execution.
|
||||
use_planning: bool,
|
||||
}
|
||||
|
||||
/// Result of a tool execution with metadata for context building.
|
||||
@@ -44,6 +48,7 @@ impl Worker {
|
||||
tools: Arc<ToolRegistry>,
|
||||
store: Option<Arc<Store>>,
|
||||
timeout: Duration,
|
||||
use_planning: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
job_id,
|
||||
@@ -53,6 +58,7 @@ impl Worker {
|
||||
tools,
|
||||
store,
|
||||
timeout,
|
||||
use_planning,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -144,6 +150,50 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
let max_iterations = 50;
|
||||
let mut iteration = 0;
|
||||
|
||||
// Generate plan if planning is enabled
|
||||
let plan = if self.use_planning {
|
||||
match reasoning.plan(reason_ctx).await {
|
||||
Ok(p) => {
|
||||
tracing::info!(
|
||||
"Created plan for job {}: {} actions, {:.0}% confidence",
|
||||
self.job_id,
|
||||
p.actions.len(),
|
||||
p.confidence * 100.0
|
||||
);
|
||||
|
||||
// Add plan to context as assistant message
|
||||
reason_ctx.messages.push(ChatMessage::assistant(format!(
|
||||
"I've created a plan to accomplish this goal: {}\n\nSteps:\n{}",
|
||||
p.goal,
|
||||
p.actions
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, a)| format!("{}. {} - {}", i + 1, a.tool_name, a.reasoning))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
)));
|
||||
|
||||
Some(p)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Planning failed for job {}, falling back to direct selection: {}",
|
||||
self.job_id,
|
||||
e
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// If we have a plan, execute it
|
||||
if let Some(ref plan) = plan {
|
||||
return self.execute_plan(rx, reasoning, reason_ctx, plan).await;
|
||||
}
|
||||
|
||||
// Otherwise, use direct tool selection loop
|
||||
loop {
|
||||
// Check for stop signal
|
||||
if let Ok(msg) = rx.try_recv() {
|
||||
@@ -401,6 +451,19 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
e
|
||||
);
|
||||
|
||||
// Record failure for self-repair tracking
|
||||
if let Some(ref store) = self.store {
|
||||
let store = store.clone();
|
||||
let tool_name = selection.tool_name.clone();
|
||||
let error_msg = e.to_string();
|
||||
tokio::spawn(async move {
|
||||
if let Err(db_err) = store.record_tool_failure(&tool_name, &error_msg).await
|
||||
{
|
||||
tracing::warn!("Failed to record tool failure: {}", db_err);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
reason_ctx.messages.push(ChatMessage::tool_result(
|
||||
"tool_call_id",
|
||||
&selection.tool_name,
|
||||
@@ -412,6 +475,95 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a pre-generated plan.
|
||||
async fn execute_plan(
|
||||
&self,
|
||||
rx: &mut mpsc::Receiver<WorkerMessage>,
|
||||
reasoning: &Reasoning,
|
||||
reason_ctx: &mut ReasoningContext,
|
||||
plan: &ActionPlan,
|
||||
) -> Result<(), Error> {
|
||||
for (i, action) in plan.actions.iter().enumerate() {
|
||||
// Check for stop signal
|
||||
if let Ok(msg) = rx.try_recv() {
|
||||
match msg {
|
||||
WorkerMessage::Stop => {
|
||||
tracing::debug!(
|
||||
"Worker for job {} received stop signal during plan execution",
|
||||
self.job_id
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
WorkerMessage::Ping => {
|
||||
tracing::trace!("Worker for job {} received ping", self.job_id);
|
||||
}
|
||||
WorkerMessage::Start => {}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
"Job {} executing planned action {}/{}: {} - {}",
|
||||
self.job_id,
|
||||
i + 1,
|
||||
plan.actions.len(),
|
||||
action.tool_name,
|
||||
action.reasoning
|
||||
);
|
||||
|
||||
// Execute the planned tool
|
||||
let result = self
|
||||
.execute_tool(&action.tool_name, &action.parameters)
|
||||
.await;
|
||||
|
||||
// Create a synthetic ToolSelection for process_tool_result
|
||||
let selection = ToolSelection {
|
||||
tool_name: action.tool_name.clone(),
|
||||
parameters: action.parameters.clone(),
|
||||
reasoning: action.reasoning.clone(),
|
||||
alternatives: vec![],
|
||||
};
|
||||
|
||||
// Process the result
|
||||
let completed = self
|
||||
.process_tool_result(reason_ctx, &selection, result)
|
||||
.await?;
|
||||
|
||||
if completed {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Small delay between actions
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
|
||||
// Plan completed, check with LLM if job is done
|
||||
reason_ctx.messages.push(ChatMessage::user(
|
||||
"All planned actions have been executed. Is the job complete? If not, what else needs to be done?",
|
||||
));
|
||||
|
||||
let response = reasoning.respond(reason_ctx).await?;
|
||||
reason_ctx.messages.push(ChatMessage::assistant(&response));
|
||||
|
||||
let response_lower = response.to_lowercase();
|
||||
if response_lower.contains("complete")
|
||||
|| response_lower.contains("finished")
|
||||
|| response_lower.contains("done")
|
||||
{
|
||||
self.mark_completed().await?;
|
||||
} else {
|
||||
// Job not complete, could re-plan or fall back to direct selection
|
||||
tracing::info!(
|
||||
"Job {} plan completed but work remains, falling back to direct selection",
|
||||
self.job_id
|
||||
);
|
||||
// Continue with standard execution loop by returning (will be picked up by main loop)
|
||||
self.mark_stuck("Plan completed but job incomplete - needs re-planning")
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn execute_tool(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
|
||||
@@ -18,6 +18,7 @@ pub struct Config {
|
||||
pub wasm: WasmConfig,
|
||||
pub secrets: SecretsConfig,
|
||||
pub builder: BuilderModeConfig,
|
||||
pub heartbeat: HeartbeatConfig,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
@@ -35,6 +36,7 @@ impl Config {
|
||||
wasm: WasmConfig::from_env()?,
|
||||
secrets: SecretsConfig::from_env()?,
|
||||
builder: BuilderModeConfig::from_env()?,
|
||||
heartbeat: HeartbeatConfig::from_env()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -166,6 +168,8 @@ pub struct AgentConfig {
|
||||
pub stuck_threshold: Duration,
|
||||
pub repair_check_interval: Duration,
|
||||
pub max_repair_attempts: u32,
|
||||
/// Whether to use planning before tool execution.
|
||||
pub use_planning: bool,
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
@@ -183,6 +187,14 @@ impl AgentConfig {
|
||||
60,
|
||||
)?),
|
||||
max_repair_attempts: parse_optional_env("SELF_REPAIR_MAX_ATTEMPTS", 3)?,
|
||||
use_planning: optional_env("AGENT_USE_PLANNING")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "AGENT_USE_PLANNING".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(true), // Default to planning enabled
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -418,6 +430,48 @@ impl BuilderModeConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Heartbeat configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HeartbeatConfig {
|
||||
/// Whether heartbeat is enabled.
|
||||
pub enabled: bool,
|
||||
/// Interval between heartbeat checks in seconds.
|
||||
pub interval_secs: u64,
|
||||
/// Channel to notify on heartbeat findings.
|
||||
pub notify_channel: Option<String>,
|
||||
/// User ID to notify on heartbeat findings.
|
||||
pub notify_user: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for HeartbeatConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
interval_secs: 1800, // 30 minutes
|
||||
notify_channel: None,
|
||||
notify_user: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HeartbeatConfig {
|
||||
fn from_env() -> Result<Self, ConfigError> {
|
||||
Ok(Self {
|
||||
enabled: optional_env("HEARTBEAT_ENABLED")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "HEARTBEAT_ENABLED".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(false),
|
||||
interval_secs: parse_optional_env("HEARTBEAT_INTERVAL_SECS", 1800)?,
|
||||
notify_channel: optional_env("HEARTBEAT_NOTIFY_CHANNEL")?,
|
||||
notify_user: optional_env("HEARTBEAT_NOTIFY_USER")?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
|
||||
fn required_env(key: &str) -> Result<String, ConfigError> {
|
||||
|
||||
@@ -439,3 +439,90 @@ fn parse_job_state(s: &str) -> JobState {
|
||||
_ => JobState::Pending,
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Tool Failures ====================
|
||||
|
||||
use crate::agent::BrokenTool;
|
||||
|
||||
impl Store {
|
||||
/// Record a tool failure (upsert: increment count if exists).
|
||||
pub async fn record_tool_failure(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
error_message: &str,
|
||||
) -> Result<(), DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
conn.execute(
|
||||
r#"
|
||||
INSERT INTO tool_failures (tool_name, error_message, error_count, last_failure)
|
||||
VALUES ($1, $2, 1, NOW())
|
||||
ON CONFLICT (tool_name) DO UPDATE SET
|
||||
error_message = $2,
|
||||
error_count = tool_failures.error_count + 1,
|
||||
last_failure = NOW()
|
||||
"#,
|
||||
&[&tool_name, &error_message],
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get tools that have failed more than `threshold` times and haven't been repaired.
|
||||
pub async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
let rows = conn
|
||||
.query(
|
||||
r#"
|
||||
SELECT tool_name, error_message, error_count, first_failure, last_failure,
|
||||
last_build_result, repair_attempts
|
||||
FROM tool_failures
|
||||
WHERE error_count >= $1 AND repaired_at IS NULL
|
||||
ORDER BY error_count DESC
|
||||
"#,
|
||||
&[&threshold],
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(rows
|
||||
.iter()
|
||||
.map(|row| BrokenTool {
|
||||
name: row.get("tool_name"),
|
||||
last_error: row.get("error_message"),
|
||||
failure_count: row.get::<_, i32>("error_count") as u32,
|
||||
first_failure: row.get("first_failure"),
|
||||
last_failure: row.get("last_failure"),
|
||||
last_build_result: row.get("last_build_result"),
|
||||
repair_attempts: row.get::<_, i32>("repair_attempts") as u32,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Mark a tool as repaired.
|
||||
pub async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
conn.execute(
|
||||
"UPDATE tool_failures SET repaired_at = NOW(), error_count = 0 WHERE tool_name = $1",
|
||||
&[&tool_name],
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Increment repair attempts for a tool.
|
||||
pub async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> {
|
||||
let conn = self.conn().await?;
|
||||
|
||||
conn.execute(
|
||||
"UPDATE tool_failures SET repair_attempts = repair_attempts + 1 WHERE tool_name = $1",
|
||||
&[&tool_name],
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
+15
-114
@@ -9,7 +9,7 @@ use std::sync::Arc;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use reqwest::Client;
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use secrecy::SecretString;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::{Mutex, RwLock};
|
||||
|
||||
@@ -127,124 +127,29 @@ impl SessionManager {
|
||||
|
||||
/// Ensure we have a valid session, triggering login flow if needed.
|
||||
///
|
||||
/// This proactively validates the token with the server, so we catch
|
||||
/// expired sessions early rather than failing on the first LLM request.
|
||||
/// If no token exists, triggers the OAuth login flow. If a token exists,
|
||||
/// it is assumed valid until a 401 response indicates otherwise.
|
||||
pub async fn ensure_authenticated(&self) -> Result<(), LlmError> {
|
||||
if !self.has_token().await {
|
||||
// No token at all, need to authenticate
|
||||
return self.initiate_login().await;
|
||||
if self.has_token().await {
|
||||
tracing::debug!("Session token present, assuming valid");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// We have a token, but let's validate it's not expired
|
||||
match self.validate_token().await {
|
||||
Ok(()) => {
|
||||
tracing::debug!("Session token validated successfully");
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Session token validation failed: {}, will re-authenticate", e);
|
||||
self.initiate_login().await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate the current token with the server.
|
||||
///
|
||||
/// Attempts to refresh the token to verify it's still valid. If refresh
|
||||
/// succeeds, we also get a fresh token as a bonus.
|
||||
async fn validate_token(&self) -> Result<(), LlmError> {
|
||||
// Try to refresh - this validates the token and gives us a fresh one
|
||||
match self.refresh_session().await {
|
||||
Ok(new_token) => {
|
||||
let mut guard = self.token.write().await;
|
||||
*guard = Some(new_token);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
// No token, need to authenticate
|
||||
self.initiate_login().await
|
||||
}
|
||||
|
||||
/// Handle an authentication failure (401 response).
|
||||
///
|
||||
/// First attempts to refresh the session. If refresh fails, initiates
|
||||
/// a full re-authentication flow.
|
||||
///
|
||||
/// Returns `true` if authentication was recovered, `false` if it failed.
|
||||
/// Triggers the OAuth login flow to get a new session token.
|
||||
pub async fn handle_auth_failure(&self) -> Result<(), LlmError> {
|
||||
// Acquire renewal lock to prevent thundering herd
|
||||
let _guard = self.renewal_lock.lock().await;
|
||||
|
||||
// Double-check: maybe another task already renewed
|
||||
// (We don't have a way to verify without making a request,
|
||||
// so we just try to refresh)
|
||||
|
||||
tracing::info!("Session expired, attempting refresh...");
|
||||
|
||||
// Try refresh first
|
||||
match self.refresh_session().await {
|
||||
Ok(new_token) => {
|
||||
let mut guard = self.token.write().await;
|
||||
*guard = Some(new_token);
|
||||
tracing::info!("Session refreshed successfully");
|
||||
return Ok(());
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Session refresh failed: {}, will need to re-authenticate",
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Refresh failed, need full re-authentication
|
||||
tracing::info!("Session expired or invalid, re-authenticating...");
|
||||
self.initiate_login().await
|
||||
}
|
||||
|
||||
/// Attempt to refresh the session using the current token.
|
||||
async fn refresh_session(&self) -> Result<SecretString, LlmError> {
|
||||
let current_token = self.get_token().await?;
|
||||
|
||||
let url = format!("{}/auth/refresh", self.config.auth_base_url);
|
||||
tracing::debug!("Attempting session refresh at {}", url);
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header(
|
||||
"Authorization",
|
||||
format!("Bearer {}", current_token.expose_secret()),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("HTTP request failed: {}", e),
|
||||
})?;
|
||||
|
||||
if response.status().is_success() {
|
||||
let body: RefreshResponse =
|
||||
response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Failed to parse response: {}", e),
|
||||
})?;
|
||||
|
||||
let new_token = SecretString::from(body.session_token.clone());
|
||||
self.save_session(&body.session_token, None).await?;
|
||||
return Ok(new_token);
|
||||
}
|
||||
|
||||
// Refresh endpoint returned non-success
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
Err(LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("HTTP {}: {}", status, body),
|
||||
})
|
||||
}
|
||||
|
||||
/// Start the OAuth login flow.
|
||||
///
|
||||
/// 1. Find an available port for the callback server
|
||||
@@ -279,8 +184,9 @@ impl SessionManager {
|
||||
})?;
|
||||
|
||||
let callback_url = format!("http://127.0.0.1:{}", port);
|
||||
// Use GitHub OAuth (Google OAuth may not have the redirect URI configured)
|
||||
let auth_url = format!(
|
||||
"{}/v1/auth/google?frontend_callback={}",
|
||||
"{}/v1/auth/github?frontend_callback={}",
|
||||
self.config.auth_base_url,
|
||||
urlencoding::encode(&callback_url)
|
||||
);
|
||||
@@ -362,8 +268,8 @@ impl SessionManager {
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
let _ = socket.shutdown().await;
|
||||
|
||||
// Provider is google since we used the google endpoint
|
||||
return Ok::<_, LlmError>((token, Some("google".to_string())));
|
||||
// Provider is github since we used the github endpoint
|
||||
return Ok::<_, LlmError>((token, Some("github".to_string())));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -480,12 +386,6 @@ impl SessionManager {
|
||||
}
|
||||
}
|
||||
|
||||
/// Response from the refresh endpoint.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RefreshResponse {
|
||||
session_token: String,
|
||||
}
|
||||
|
||||
/// Create a session manager from a config, migrating from env var if present.
|
||||
pub async fn create_session_manager(config: SessionConfig) -> Arc<SessionManager> {
|
||||
let manager = SessionManager::new_async(config).await;
|
||||
@@ -509,6 +409,7 @@ pub async fn create_session_manager(config: SessionConfig) -> Arc<SessionManager
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use secrecy::ExposeSecret;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -192,6 +192,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
tools,
|
||||
channels,
|
||||
workspace,
|
||||
Some(config.heartbeat.clone()),
|
||||
);
|
||||
|
||||
tracing::info!("Agent initialized, starting main loop...");
|
||||
|
||||
Reference in New Issue
Block a user