feat: 10 infrastructure improvements from zeroclaw (#126)

* refactor: break up agent_loop.rs into four focused modules

Split the monolithic 2835-line agent_loop.rs into:
- agent_loop.rs (722L): Agent struct, event loop, message dispatch
- dispatcher.rs (635L): Agentic tool loop, tool execution, auth detection
- commands.rs (484L): System commands, job handlers, heartbeat, summarize
- thread_ops.rs (1059L): Thread lifecycle, approval, undo/redo, persistence

Each module gets its own impl Agent block. Agent fields changed to
pub(super) so sibling modules in the agent package can access them.
All 16 existing tests pass in their new locations.

Inspired by ZeroClaw's agent module split (agent.rs, loop_.rs,
dispatcher.rs, prompt.rs, memory_loader.rs).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add cost caps and guardrails for autonomous agent spending

Daily budget (MAX_COST_PER_DAY_CENTS) and hourly action rate
(MAX_ACTIONS_PER_HOUR) limits prevent runaway agents from burning
through API credits, especially in daemon/heartbeat modes.

- CostGuard with pre-flight check and post-call recording
- Sliding window for hourly rate, midnight-UTC daily reset
- 80% threshold warning, atomic fast-path for exceeded budget
- Wired into dispatcher loop (check before LLM call, record after)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add circuit breaker on LLM providers

Wraps LlmProvider with a Closed/Open/HalfOpen state machine that
trips after consecutive transient failures, preventing request storms
against a degraded backend. Automatically probes for recovery.

- CircuitBreakerProvider implements LlmProvider (drop-in wrapper)
- Transient error classification (server, rate-limit, network, auth infra)
- Client errors (wrong model, context overflow) don't trip the breaker
- Configurable via CIRCUIT_BREAKER_THRESHOLD and CIRCUIT_BREAKER_RECOVERY_SECS
- Composes with existing FailoverProvider (circuit breaker wraps failover)
- 12 tests covering full state machine and error classification

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add tunnel abstraction for remote access

Trait-based tunnel system with lifecycle management (start/stop/health)
for exposing the agent to the internet through external tunnel binaries.

Five providers:
- Cloudflare Tunnel (cloudflared, Zero Trust token auth)
- Tailscale (serve for tailnet, funnel for public)
- ngrok (with optional custom domain)
- Custom (arbitrary command with {host}/{port} placeholders)
- None (local-only, no external exposure)

Config via TUNNEL_PROVIDER + provider-specific env vars. Extends
existing TunnelConfig with optional managed provider alongside the
static TUNNEL_URL path. Factory, shared process management, and
37 tests covering all providers and edge cases.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add OS service management (launchd/systemd)

Adds `ironclaw service {install,start,stop,status,uninstall}` for
running the agent as a background daemon. macOS uses launchd plists
under ~/Library/LaunchAgents, Linux uses systemd user units.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add observability trait system with noop, log, and multi backends

Introduces an Observer trait for recording agent lifecycle events and
metrics, with pluggable backends. The noop backend compiles to zero
overhead, log backend uses tracing, and multi fans out to multiple
observers. Configured via OBSERVABILITY_BACKEND env var.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add in-memory LLM response cache with TTL and LRU eviction

CachedProvider wraps any LlmProvider and caches complete() responses
keyed by SHA-256(model + messages). Tool-calling requests are never
cached since they trigger side effects. Configurable via
RESPONSE_CACHE_ENABLED, RESPONSE_CACHE_TTL_SECS, and
RESPONSE_CACHE_MAX_ENTRIES env vars.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add memory hygiene with cadence-gated daily log cleanup

Adds workspace::hygiene module that automatically deletes daily log
documents older than a configurable retention period (default 30 days).
Runs on a 12-hour cadence tracked via a local state file to avoid
redundant passes. Best-effort design: failures are logged, never fatal.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add doctor diagnostics command for active health probing

Probes external dependencies (Docker, cloudflared, ngrok, tailscale),
validates NEAR AI session, checks database connectivity, and verifies
workspace directory. Complements the passive `status` command.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add structured TOML config file support

Adds ~/.ironclaw/config.toml as a configuration layer between env vars
and database settings. Priority: env var > TOML file > DB > defaults.

- `ironclaw config init` generates a commented config.toml from current settings
- `ironclaw --config path/to/config.toml` loads a custom config file
- Settings.merge_from() only overlays non-default values from the TOML file
- `ironclaw config path` now shows TOML file status

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address codex review findings

- apply_toml_overlay now returns Result and errors on explicit missing
  or invalid config paths (was log-only, violating the documented
  contract that explicit paths are fatal)
- custom tunnel url_pattern is now used to filter extracted URLs, not
  just as a gate for scanning stdout
- systemd ExecStart path is now quoted to handle spaces in paths

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address PR review feedback

- Cache key now includes max_tokens, temperature, and stop_sequences
  so different request parameters produce distinct keys
- to_cents() uses .trunc() + parse::<u64> instead of f64 intermediary,
  avoiding precision loss for large values
- Tailscale public URL no longer includes local port (serve/funnel
  expose on standard HTTPS port 443)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: wire up tunnel lifecycle and fix audit findings

Connect the tunnel module to the rest of the application so that
setting TUNNEL_PROVIDER actually starts a managed tunnel at boot and
stops it on shutdown. Previously create_tunnel() was never called
outside tests.

Changes:
- Expand TunnelSettings with provider credential fields (settings.rs)
- TunnelConfig::resolve() falls back to DB settings when env vars unset
- Start tunnel at boot, stop on shutdown, show URL in boot screen
- Setup wizard collects provider-specific credentials (ngrok, cloudflare,
  tailscale, custom, static URL)
- Fix public_url() returning None under lock contention (SharedUrl)
- Fix local_host parameter ignored by cloudflare/ngrok/tailscale
- Fix tailscale silent fallback to "localhost" on bad JSON
- Fix ngrok globally mutating config via add-authtoken (use env var)
- Add 10s timeout to tailscale status --json

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address PR review comments

- Document split_whitespace limitation in CustomTunnel doc comment
- Remove unnecessary quotes from systemd ExecStart directive

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address PR review feedback (round 3)

- doctor: missing libSQL DB on fresh install is Pass, not Fail
- service: quote ExecStart path for systemd space handling

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: correct cost guard doc comment (LLM calls, not LLM/tool)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-17 19:35:56 +00:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 436dda0f2f
commit a158eee1b0
37 changed files with 7272 additions and 2177 deletions
Generated
+1
View File
@@ -2544,6 +2544,7 @@ dependencies = [
"tokio-stream",
"tokio-test",
"tokio-tungstenite 0.26.2",
"toml",
"tower 0.5.3",
"tower-http 0.6.8",
"tracing",
+1
View File
@@ -55,6 +55,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
# Configuration
dotenvy = "0.15"
toml = "0.8"
# Core types
uuid = { version = "1", features = ["v4", "serde"] }
+33 -2127
View File
File diff suppressed because it is too large Load Diff
+484
View File
@@ -0,0 +1,484 @@
//! System commands and job handlers for the agent.
//!
//! Extracted from `agent_loop.rs` to isolate the /help, /model, /status,
//! and other command processing from the core agent loop.
use std::sync::Arc;
use tokio::sync::Mutex;
use uuid::Uuid;
use crate::agent::session::Session;
use crate::agent::submission::SubmissionResult;
use crate::agent::{Agent, MessageIntent};
use crate::channels::{IncomingMessage, StatusUpdate};
use crate::error::Error;
use crate::llm::ChatMessage;
impl Agent {
/// Handle job-related intents without turn tracking.
pub(super) async fn handle_job_or_command(
&self,
intent: MessageIntent,
message: &IncomingMessage,
) -> Result<SubmissionResult, Error> {
// Send thinking status for non-trivial operations
if let MessageIntent::CreateJob { .. } = &intent {
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::Thinking("Processing...".into()),
&message.metadata,
)
.await;
}
let response = match intent {
MessageIntent::CreateJob {
title,
description,
category,
} => {
self.handle_create_job(&message.user_id, title, description, category)
.await?
}
MessageIntent::CheckJobStatus { job_id } => {
self.handle_check_status(&message.user_id, job_id).await?
}
MessageIntent::CancelJob { job_id } => {
self.handle_cancel_job(&message.user_id, &job_id).await?
}
MessageIntent::ListJobs { filter } => {
self.handle_list_jobs(&message.user_id, filter).await?
}
MessageIntent::HelpJob { job_id } => {
self.handle_help_job(&message.user_id, &job_id).await?
}
MessageIntent::Command { command, args } => {
match self.handle_command(&command, &args).await? {
Some(s) => s,
None => return Ok(SubmissionResult::Ok { message: None }), // Shutdown signal
}
}
_ => "Unknown intent".to_string(),
};
Ok(SubmissionResult::response(response))
}
async fn handle_create_job(
&self,
user_id: &str,
title: String,
description: String,
category: Option<String>,
) -> Result<String, Error> {
// Create job context
let job_id = self
.context_manager
.create_job_for_user(user_id, &title, &description)
.await?;
// Update category if provided
if let Some(cat) = category {
self.context_manager
.update_context(job_id, |ctx| {
ctx.category = Some(cat);
})
.await?;
}
// Persist new job to database (fire-and-forget)
if let Some(store) = self.store()
&& let Ok(ctx) = self.context_manager.get_context(job_id).await
{
let store = store.clone();
tokio::spawn(async move {
if let Err(e) = store.save_job(&ctx).await {
tracing::warn!("Failed to persist new job {}: {}", job_id, e);
}
});
}
// Schedule for execution
self.scheduler.schedule(job_id).await?;
Ok(format!(
"Created job: {}\nID: {}\n\nThe job has been scheduled and is now running.",
title, job_id
))
}
async fn handle_check_status(
&self,
user_id: &str,
job_id: Option<String>,
) -> Result<String, Error> {
match job_id {
Some(id) => {
let uuid = Uuid::parse_str(&id)
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
let ctx = self.context_manager.get_context(uuid).await?;
if ctx.user_id != user_id {
return Err(crate::error::JobError::NotFound { id: uuid }.into());
}
Ok(format!(
"Job: {}\nStatus: {:?}\nCreated: {}\nStarted: {}\nActual cost: {}",
ctx.title,
ctx.state,
ctx.created_at.format("%Y-%m-%d %H:%M:%S"),
ctx.started_at
.map(|t| t.format("%Y-%m-%d %H:%M:%S").to_string())
.unwrap_or_else(|| "Not started".to_string()),
ctx.actual_cost
))
}
None => {
// Show summary of all jobs
let summary = self.context_manager.summary_for(user_id).await;
Ok(format!(
"Jobs summary:\n Total: {}\n In Progress: {}\n Completed: {}\n Failed: {}\n Stuck: {}",
summary.total,
summary.in_progress,
summary.completed,
summary.failed,
summary.stuck
))
}
}
}
async fn handle_cancel_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
let uuid = Uuid::parse_str(job_id)
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
let ctx = self.context_manager.get_context(uuid).await?;
if ctx.user_id != user_id {
return Err(crate::error::JobError::NotFound { id: uuid }.into());
}
self.scheduler.stop(uuid).await?;
Ok(format!("Job {} has been cancelled.", job_id))
}
async fn handle_list_jobs(
&self,
user_id: &str,
_filter: Option<String>,
) -> Result<String, Error> {
let jobs = self.context_manager.all_jobs_for(user_id).await;
if jobs.is_empty() {
return Ok("No jobs found.".to_string());
}
let mut output = String::from("Jobs:\n");
for job_id in jobs {
if let Ok(ctx) = self.context_manager.get_context(job_id).await
&& ctx.user_id == user_id
{
output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state));
}
}
Ok(output)
}
async fn handle_help_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
let uuid = Uuid::parse_str(job_id)
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
let ctx = self.context_manager.get_context(uuid).await?;
if ctx.user_id != user_id {
return Err(crate::error::JobError::NotFound { id: uuid }.into());
}
if ctx.state == crate::context::JobState::Stuck {
// Attempt recovery
self.context_manager
.update_context(uuid, |ctx| ctx.attempt_recovery())
.await?
.map_err(|s| crate::error::JobError::ContextError {
id: uuid,
reason: s,
})?;
// Reschedule
self.scheduler.schedule(uuid).await?;
Ok(format!(
"Job {} was stuck. Attempting recovery (attempt #{}).",
job_id,
ctx.repair_attempts + 1
))
} else {
Ok(format!(
"Job {} is not stuck (current state: {:?}). No help needed.",
job_id, ctx.state
))
}
}
/// Trigger a manual heartbeat check.
pub(super) async fn process_heartbeat(&self) -> Result<SubmissionResult, Error> {
let Some(workspace) = self.workspace() else {
return Ok(SubmissionResult::error(
"Heartbeat requires a workspace (database must be connected).",
));
};
let runner = crate::agent::HeartbeatRunner::new(
crate::agent::HeartbeatConfig::default(),
workspace.clone(),
self.llm().clone(),
);
match runner.check_heartbeat().await {
crate::agent::HeartbeatResult::Ok => Ok(SubmissionResult::ok_with_message(
"Heartbeat: all clear, nothing needs attention.",
)),
crate::agent::HeartbeatResult::NeedsAttention(msg) => Ok(SubmissionResult::response(
format!("Heartbeat findings:\n\n{}", msg),
)),
crate::agent::HeartbeatResult::Skipped => Ok(SubmissionResult::ok_with_message(
"Heartbeat skipped: no HEARTBEAT.md checklist found in workspace.",
)),
crate::agent::HeartbeatResult::Failed(err) => Ok(SubmissionResult::error(format!(
"Heartbeat failed: {}",
err
))),
}
}
/// Summarize the current thread's conversation.
pub(super) async fn process_summarize(
&self,
session: Arc<Mutex<Session>>,
thread_id: Uuid,
) -> Result<SubmissionResult, Error> {
let messages = {
let sess = session.lock().await;
let thread = sess
.threads
.get(&thread_id)
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
thread.messages()
};
if messages.is_empty() {
return Ok(SubmissionResult::ok_with_message(
"Nothing to summarize (empty thread).",
));
}
// Build a summary prompt with the conversation
let mut context = Vec::new();
context.push(ChatMessage::system(
"Summarize the conversation so far in 3-5 concise bullet points. \
Focus on decisions made, actions taken, and key outcomes. \
Be brief and factual.",
));
// Include the conversation messages (truncate to last 20 to avoid context overflow)
let start = if messages.len() > 20 {
messages.len() - 20
} else {
0
};
context.extend_from_slice(&messages[start..]);
context.push(ChatMessage::user("Summarize this conversation."));
let request = crate::llm::CompletionRequest::new(context)
.with_max_tokens(512)
.with_temperature(0.3);
match self.llm().complete(request).await {
Ok(response) => Ok(SubmissionResult::response(format!(
"Thread Summary:\n\n{}",
response.content.trim()
))),
Err(e) => Ok(SubmissionResult::error(format!("Summarize failed: {}", e))),
}
}
/// Suggest next steps based on the current thread.
pub(super) async fn process_suggest(
&self,
session: Arc<Mutex<Session>>,
thread_id: Uuid,
) -> Result<SubmissionResult, Error> {
let messages = {
let sess = session.lock().await;
let thread = sess
.threads
.get(&thread_id)
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
thread.messages()
};
if messages.is_empty() {
return Ok(SubmissionResult::ok_with_message(
"Nothing to suggest from (empty thread).",
));
}
let mut context = Vec::new();
context.push(ChatMessage::system(
"Based on the conversation so far, suggest 2-4 concrete next steps the user could take. \
Be actionable and specific. Format as a numbered list.",
));
let start = if messages.len() > 20 {
messages.len() - 20
} else {
0
};
context.extend_from_slice(&messages[start..]);
context.push(ChatMessage::user("What should I do next?"));
let request = crate::llm::CompletionRequest::new(context)
.with_max_tokens(512)
.with_temperature(0.5);
match self.llm().complete(request).await {
Ok(response) => Ok(SubmissionResult::response(format!(
"Suggested Next Steps:\n\n{}",
response.content.trim()
))),
Err(e) => Ok(SubmissionResult::error(format!("Suggest failed: {}", e))),
}
}
/// Handle system commands that bypass thread-state checks entirely.
pub(super) async fn handle_system_command(
&self,
command: &str,
args: &[String],
) -> Result<SubmissionResult, Error> {
match command {
"help" => Ok(SubmissionResult::response(concat!(
"System:\n",
" /help Show this help\n",
" /model [name] Show or switch the active model\n",
" /version Show version info\n",
" /tools List available tools\n",
" /debug Toggle debug mode\n",
" /ping Connectivity check\n",
"\n",
"Jobs:\n",
" /job <desc> Create a new job\n",
" /status [id] Check job status\n",
" /cancel <id> Cancel a job\n",
" /list List all jobs\n",
"\n",
"Session:\n",
" /undo Undo last turn\n",
" /redo Redo undone turn\n",
" /compact Compress context window\n",
" /clear Clear current thread\n",
" /interrupt Stop current operation\n",
" /new New conversation thread\n",
" /thread <id> Switch to thread\n",
" /resume <id> Resume from checkpoint\n",
"\n",
"Agent:\n",
" /heartbeat Run heartbeat check\n",
" /summarize Summarize current thread\n",
" /suggest Suggest next steps\n",
"\n",
" /quit Exit",
))),
"ping" => Ok(SubmissionResult::response("pong!")),
"version" => Ok(SubmissionResult::response(format!(
"{} v{}",
env!("CARGO_PKG_NAME"),
env!("CARGO_PKG_VERSION")
))),
"tools" => {
let tools = self.tools().list().await;
Ok(SubmissionResult::response(format!(
"Available tools: {}",
tools.join(", ")
)))
}
"debug" => {
// Debug toggle is handled client-side in the REPL.
// For non-REPL channels, just acknowledge.
Ok(SubmissionResult::ok_with_message(
"Debug toggle is handled by your client.",
))
}
"model" => {
if args.is_empty() {
// Show current model
let name = self.llm().active_model_name();
Ok(SubmissionResult::response(format!(
"Active model: {}",
name
)))
} else {
let requested = &args[0];
// Validate the model exists
match self.llm().list_models().await {
Ok(models) if !models.is_empty() => {
if !models.iter().any(|m| m == requested) {
return Ok(SubmissionResult::error(format!(
"Unknown model: {}. Available models:\n {}",
requested,
models.join("\n ")
)));
}
}
Ok(_) => {
// Empty model list, can't validate but try anyway
}
Err(e) => {
tracing::warn!("Could not fetch model list for validation: {}", e);
// Proceed anyway, the provider will error on the next call if invalid
}
}
match self.llm().set_model(requested) {
Ok(()) => Ok(SubmissionResult::response(format!(
"Switched model to: {}",
requested
))),
Err(e) => Ok(SubmissionResult::error(format!(
"Failed to switch model: {}",
e
))),
}
}
}
_ => Ok(SubmissionResult::error(format!(
"Unknown command: {}. Try /help",
command
))),
}
}
/// Handle legacy command routing from the Router (job commands that go through
/// process_user_input -> router -> handle_job_or_command -> here).
pub(super) async fn handle_command(
&self,
command: &str,
args: &[String],
) -> Result<Option<String>, Error> {
// System commands are now handled directly via Submission::SystemCommand,
// but the router may still send us unknown /commands.
match self.handle_system_command(command, args).await? {
SubmissionResult::Response { content } => Ok(Some(content)),
SubmissionResult::Ok { message } => Ok(message),
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
_ => Ok(None),
}
}
}
+339
View File
@@ -0,0 +1,339 @@
//! Cost enforcement guardrails for the agent.
//!
//! Tracks LLM spending and action rates, enforcing configurable limits
//! to prevent runaway agents from burning through API credits. Especially
//! important for daemon/heartbeat modes where the agent acts autonomously.
use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Instant;
use rust_decimal::Decimal;
use rust_decimal_macros::dec;
use tokio::sync::Mutex;
use crate::llm::costs;
/// Configuration for cost guardrails.
#[derive(Debug, Clone, Default)]
pub struct CostGuardConfig {
/// Maximum spend per day in cents (e.g. 10000 = $100). None = unlimited.
pub max_cost_per_day_cents: Option<u64>,
/// Maximum LLM calls per hour. None = unlimited.
pub max_actions_per_hour: Option<u64>,
}
/// Error returned when a cost limit is exceeded.
#[derive(Debug, Clone)]
pub enum CostLimitExceeded {
/// Daily spending cap reached.
DailyBudget { spent_cents: u64, limit_cents: u64 },
/// Hourly action rate limit reached.
HourlyRate { actions: u64, limit: u64 },
}
impl std::fmt::Display for CostLimitExceeded {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::DailyBudget {
spent_cents,
limit_cents,
} => write!(
f,
"Daily cost limit exceeded: spent ${:.2} of ${:.2} allowed",
*spent_cents as f64 / 100.0,
*limit_cents as f64 / 100.0
),
Self::HourlyRate { actions, limit } => write!(
f,
"Hourly action limit exceeded: {} actions of {} allowed per hour",
actions, limit
),
}
}
}
/// Tracks costs and action rates, enforcing configurable limits.
///
/// Thread-safe; designed to be shared via `Arc<CostGuard>`.
pub struct CostGuard {
config: CostGuardConfig,
/// Running cost total for the current day (in USD, not cents).
daily_cost: Mutex<DailyCost>,
/// Sliding window of action timestamps for rate limiting.
action_window: Mutex<VecDeque<Instant>>,
/// Flag set when daily budget is exceeded to short-circuit checks.
budget_exceeded: AtomicBool,
}
struct DailyCost {
total: Decimal,
/// Day boundary (midnight UTC) for resetting the counter.
reset_date: chrono::NaiveDate,
}
impl CostGuard {
pub fn new(config: CostGuardConfig) -> Self {
Self {
config,
daily_cost: Mutex::new(DailyCost {
total: Decimal::ZERO,
reset_date: chrono::Utc::now().date_naive(),
}),
action_window: Mutex::new(VecDeque::new()),
budget_exceeded: AtomicBool::new(false),
}
}
/// Check whether the next action is allowed under the configured limits.
///
/// Call this BEFORE making an LLM call. Does NOT record the action yet,
/// call `record_action` after the action completes.
pub async fn check_allowed(&self) -> Result<(), CostLimitExceeded> {
// Fast path: if budget already blown, skip the lock
if self.budget_exceeded.load(Ordering::Relaxed) {
let daily = self.daily_cost.lock().await;
let spent_cents = to_cents(daily.total);
return Err(CostLimitExceeded::DailyBudget {
spent_cents,
limit_cents: self.config.max_cost_per_day_cents.unwrap_or(0),
});
}
// Check daily budget
if let Some(limit_cents) = self.config.max_cost_per_day_cents {
let daily = self.daily_cost.lock().await;
let spent_cents = to_cents(daily.total);
if spent_cents >= limit_cents {
self.budget_exceeded.store(true, Ordering::Relaxed);
return Err(CostLimitExceeded::DailyBudget {
spent_cents,
limit_cents,
});
}
}
// Check hourly rate
if let Some(limit) = self.config.max_actions_per_hour {
let mut window = self.action_window.lock().await;
let cutoff = Instant::now() - std::time::Duration::from_secs(3600);
// Drain expired entries
while window.front().is_some_and(|t| *t < cutoff) {
window.pop_front();
}
let count = window.len() as u64;
if count >= limit {
return Err(CostLimitExceeded::HourlyRate {
actions: count,
limit,
});
}
}
Ok(())
}
/// Record a completed LLM action: its token costs and the action timestamp.
///
/// Call this AFTER an LLM call completes so that costs are tracked.
pub async fn record_llm_call(
&self,
model: &str,
input_tokens: u32,
output_tokens: u32,
) -> Decimal {
let (input_rate, output_rate) =
costs::model_cost(model).unwrap_or_else(costs::default_cost);
let cost =
input_rate * Decimal::from(input_tokens) + output_rate * Decimal::from(output_tokens);
// Update daily cost (reset if new day)
{
let mut daily = self.daily_cost.lock().await;
let today = chrono::Utc::now().date_naive();
if today != daily.reset_date {
daily.total = Decimal::ZERO;
daily.reset_date = today;
self.budget_exceeded.store(false, Ordering::Relaxed);
tracing::info!("Cost guard: daily counter reset for {}", today);
}
daily.total += cost;
// Check if we just crossed the threshold
if let Some(limit_cents) = self.config.max_cost_per_day_cents {
let spent_cents = to_cents(daily.total);
if spent_cents >= limit_cents {
self.budget_exceeded.store(true, Ordering::Relaxed);
tracing::warn!(
"Daily cost limit reached: ${:.2} of ${:.2}",
daily.total,
Decimal::from(limit_cents) / dec!(100)
);
}
// Warn at 80% threshold
let warn_threshold = limit_cents * 80 / 100;
if spent_cents >= warn_threshold && spent_cents < limit_cents {
tracing::warn!(
"Approaching daily cost limit: ${:.2} of ${:.2} ({}%)",
daily.total,
Decimal::from(limit_cents) / dec!(100),
spent_cents * 100 / limit_cents
);
}
}
}
// Record action in sliding window
{
let mut window = self.action_window.lock().await;
window.push_back(Instant::now());
}
cost
}
/// Current daily spend in USD (as Decimal).
pub async fn daily_spend(&self) -> Decimal {
let daily = self.daily_cost.lock().await;
let today = chrono::Utc::now().date_naive();
if today != daily.reset_date {
Decimal::ZERO
} else {
daily.total
}
}
/// Number of actions in the current hourly window.
pub async fn actions_this_hour(&self) -> u64 {
let mut window = self.action_window.lock().await;
let cutoff = Instant::now() - std::time::Duration::from_secs(3600);
while window.front().is_some_and(|t| *t < cutoff) {
window.pop_front();
}
window.len() as u64
}
}
/// Convert a Decimal USD amount to whole cents (truncated).
fn to_cents(usd: Decimal) -> u64 {
let cents = (usd * dec!(100)).trunc();
cents.to_string().parse::<u64>().unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_unlimited_allows_everything() {
let guard = CostGuard::new(CostGuardConfig::default());
// No limits set, should always be allowed
assert!(guard.check_allowed().await.is_ok());
// Record a big call, still allowed
guard.record_llm_call("gpt-4o", 100_000, 100_000).await;
assert!(guard.check_allowed().await.is_ok());
}
#[tokio::test]
async fn test_daily_budget_enforcement() {
let guard = CostGuard::new(CostGuardConfig {
max_cost_per_day_cents: Some(1), // $0.01 limit
max_actions_per_hour: None,
});
// First call allowed
assert!(guard.check_allowed().await.is_ok());
// Record a call that costs more than $0.01
// gpt-4o: input=$0.0000025/tok, output=$0.00001/tok
// 10000 input + 10000 output = $0.025 + $0.10 = $0.125
guard.record_llm_call("gpt-4o", 10_000, 10_000).await;
// Now should be blocked
let result = guard.check_allowed().await;
assert!(result.is_err());
match result.unwrap_err() {
CostLimitExceeded::DailyBudget { limit_cents, .. } => {
assert_eq!(limit_cents, 1);
}
other => panic!("Expected DailyBudget, got {:?}", other),
}
}
#[tokio::test]
async fn test_hourly_rate_enforcement() {
let guard = CostGuard::new(CostGuardConfig {
max_cost_per_day_cents: None,
max_actions_per_hour: Some(3),
});
// First 3 actions allowed
for _ in 0..3 {
assert!(guard.check_allowed().await.is_ok());
guard.record_llm_call("gpt-4o", 10, 10).await;
}
// 4th should be blocked
let result = guard.check_allowed().await;
assert!(result.is_err());
match result.unwrap_err() {
CostLimitExceeded::HourlyRate { actions, limit } => {
assert_eq!(actions, 3);
assert_eq!(limit, 3);
}
other => panic!("Expected HourlyRate, got {:?}", other),
}
}
#[tokio::test]
async fn test_daily_spend_tracking() {
let guard = CostGuard::new(CostGuardConfig::default());
assert_eq!(guard.daily_spend().await, Decimal::ZERO);
let cost = guard.record_llm_call("gpt-4o", 1000, 500).await;
assert!(cost > Decimal::ZERO);
assert_eq!(guard.daily_spend().await, cost);
}
#[tokio::test]
async fn test_actions_this_hour() {
let guard = CostGuard::new(CostGuardConfig::default());
assert_eq!(guard.actions_this_hour().await, 0);
guard.record_llm_call("gpt-4o", 10, 10).await;
guard.record_llm_call("gpt-4o", 10, 10).await;
assert_eq!(guard.actions_this_hour().await, 2);
}
#[test]
fn test_to_cents() {
assert_eq!(to_cents(dec!(1.50)), 150);
assert_eq!(to_cents(dec!(0.01)), 1);
assert_eq!(to_cents(Decimal::ZERO), 0);
}
#[test]
fn test_cost_limit_display() {
let budget = CostLimitExceeded::DailyBudget {
spent_cents: 1050,
limit_cents: 1000,
};
assert!(budget.to_string().contains("$10.50"));
assert!(budget.to_string().contains("$10.00"));
let rate = CostLimitExceeded::HourlyRate {
actions: 101,
limit: 100,
};
assert!(rate.to_string().contains("101 actions"));
assert!(rate.to_string().contains("100 allowed"));
}
}
+635
View File
@@ -0,0 +1,635 @@
//! Tool dispatch logic for the agent.
//!
//! Extracted from `agent_loop.rs` to keep the core agentic tool execution
//! loop (LLM call -> tool calls -> repeat) in its own focused module.
use std::sync::Arc;
use tokio::sync::Mutex;
use uuid::Uuid;
use crate::agent::Agent;
use crate::agent::session::{PendingApproval, Session, ThreadState};
use crate::channels::{IncomingMessage, StatusUpdate};
use crate::context::JobContext;
use crate::error::Error;
use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult};
/// Result of the agentic loop execution.
pub(super) enum AgenticLoopResult {
/// Completed with a response.
Response(String),
/// A tool requires approval before continuing.
NeedApproval {
/// The pending approval request to store.
pending: PendingApproval,
},
}
impl Agent {
/// Run the agentic loop: call LLM, execute tools, repeat until text response.
///
/// Returns `AgenticLoopResult::Response` on completion, or
/// `AgenticLoopResult::NeedApproval` if a tool requires user approval.
///
/// When `resume_after_tool` is true the loop already knows a tool was
/// executed earlier in this turn (e.g. an approved tool), so it won't
/// force the LLM to use tools if it responds with text.
pub(super) async fn run_agentic_loop(
&self,
message: &IncomingMessage,
session: Arc<Mutex<Session>>,
thread_id: Uuid,
initial_messages: Vec<ChatMessage>,
resume_after_tool: bool,
) -> Result<AgenticLoopResult, Error> {
// Load workspace system prompt (identity files: AGENTS.md, SOUL.md, etc.)
let system_prompt = if let Some(ws) = self.workspace() {
match ws.system_prompt().await {
Ok(prompt) if !prompt.is_empty() => Some(prompt),
Ok(_) => None,
Err(e) => {
tracing::debug!("Could not load workspace system prompt: {}", e);
None
}
}
} else {
None
};
let mut reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
if let Some(prompt) = system_prompt {
reasoning = reasoning.with_system_prompt(prompt);
}
// Build context with messages that we'll mutate during the loop
let mut context_messages = initial_messages;
// Create a JobContext for tool execution (chat doesn't have a real job)
let job_ctx = JobContext::with_user(&message.user_id, "chat", "Interactive chat session");
const MAX_TOOL_ITERATIONS: usize = 10;
let mut iteration = 0;
let mut tools_executed = resume_after_tool;
loop {
iteration += 1;
if iteration > MAX_TOOL_ITERATIONS {
return Err(crate::error::LlmError::InvalidResponse {
provider: "agent".to_string(),
reason: format!("Exceeded maximum tool iterations ({})", MAX_TOOL_ITERATIONS),
}
.into());
}
// Check if interrupted
{
let sess = session.lock().await;
if let Some(thread) = sess.threads.get(&thread_id)
&& thread.state == ThreadState::Interrupted
{
return Err(crate::error::JobError::ContextError {
id: thread_id,
reason: "Interrupted".to_string(),
}
.into());
}
}
// Enforce cost guardrails before the LLM call
if let Err(limit) = self.cost_guard().check_allowed().await {
return Err(crate::error::LlmError::InvalidResponse {
provider: "agent".to_string(),
reason: limit.to_string(),
}
.into());
}
// Refresh tool definitions each iteration so newly built tools become visible
let tool_defs = self.tools().tool_definitions().await;
// Call LLM with current context
let context = ReasoningContext::new()
.with_messages(context_messages.clone())
.with_tools(tool_defs)
.with_metadata({
let mut m = std::collections::HashMap::new();
m.insert("thread_id".to_string(), thread_id.to_string());
m
});
let output = reasoning.respond_with_tools(&context).await?;
// Record cost and track token usage
let model_name = self.llm().active_model_name();
let call_cost = self
.cost_guard()
.record_llm_call(
&model_name,
output.usage.input_tokens,
output.usage.output_tokens,
)
.await;
tracing::debug!(
"LLM call used {} input + {} output tokens (${:.6})",
output.usage.input_tokens,
output.usage.output_tokens,
call_cost,
);
match output.result {
RespondResult::Text(text) => {
// If no tools have been executed yet, prompt the LLM to use tools
// This handles the case where the model explains what it will do
// instead of actually calling tools
if !tools_executed && iteration < 3 {
tracing::debug!(
"No tools executed yet (iteration {}), prompting for tool use",
iteration
);
context_messages.push(ChatMessage::assistant(&text));
context_messages.push(ChatMessage::user(
"Please proceed and use the available tools to complete this task.",
));
continue;
}
// Tools have been executed or we've tried multiple times, return response
return Ok(AgenticLoopResult::Response(text));
}
RespondResult::ToolCalls {
tool_calls,
content,
} => {
tools_executed = true;
// Add the assistant message with tool_calls to context.
// OpenAI protocol requires this before tool-result messages.
context_messages.push(ChatMessage::assistant_with_tool_calls(
content,
tool_calls.clone(),
));
// Execute tools and add results to context
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::Thinking(format!(
"Executing {} tool(s)...",
tool_calls.len()
)),
&message.metadata,
)
.await;
// Record tool calls in the thread
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
for tc in &tool_calls {
turn.record_tool_call(&tc.name, tc.arguments.clone());
}
}
}
// Execute each tool (with approval checking and hook interception)
for mut tc in tool_calls {
// Check if tool requires approval
if let Some(tool) = self.tools().get(&tc.name).await
&& tool.requires_approval()
{
// Check if auto-approved for this session
let mut is_auto_approved = {
let sess = session.lock().await;
sess.is_tool_auto_approved(&tc.name)
};
// Override auto-approval for destructive parameters
// (e.g. `rm -rf`, `git push --force` in shell commands).
if is_auto_approved && tool.requires_approval_for(&tc.arguments) {
tracing::info!(
tool = %tc.name,
"Parameters require explicit approval despite auto-approve"
);
is_auto_approved = false;
}
if !is_auto_approved {
// Need approval - store pending request and return
let pending = PendingApproval {
request_id: Uuid::new_v4(),
tool_name: tc.name.clone(),
parameters: tc.arguments.clone(),
description: tool.description().to_string(),
tool_call_id: tc.id.clone(),
context_messages: context_messages.clone(),
};
return Ok(AgenticLoopResult::NeedApproval { pending });
}
}
// Hook: BeforeToolCall — allow hooks to modify or reject tool calls
{
let event = crate::hooks::HookEvent::ToolCall {
tool_name: tc.name.clone(),
parameters: tc.arguments.clone(),
user_id: message.user_id.clone(),
context: "chat".to_string(),
};
match self.hooks().run(&event).await {
Err(crate::hooks::HookError::Rejected { reason }) => {
context_messages.push(ChatMessage::tool_result(
&tc.id,
&tc.name,
format!("Tool call rejected by hook: {}", reason),
));
continue;
}
Err(err) => {
context_messages.push(ChatMessage::tool_result(
&tc.id,
&tc.name,
format!("Tool call blocked by hook policy: {}", err),
));
continue;
}
Ok(crate::hooks::HookOutcome::Continue {
modified: Some(new_params),
}) => match serde_json::from_str(&new_params) {
Ok(parsed) => tc.arguments = parsed,
Err(e) => {
tracing::warn!(
tool = %tc.name,
"Hook returned non-JSON modification for ToolCall, ignoring: {}",
e
);
}
},
_ => {} // Continue, fail-open errors already logged
}
}
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::ToolStarted {
name: tc.name.clone(),
},
&message.metadata,
)
.await;
let tool_result = self
.execute_chat_tool(&tc.name, &tc.arguments, &job_ctx)
.await;
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::ToolCompleted {
name: tc.name.clone(),
success: tool_result.is_ok(),
},
&message.metadata,
)
.await;
if let Ok(ref output) = tool_result
&& !output.is_empty()
{
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::ToolResult {
name: tc.name.clone(),
preview: output.clone(),
},
&message.metadata,
)
.await;
}
// Record result in thread
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
match &tool_result {
Ok(output) => {
turn.record_tool_result(serde_json::json!(output));
}
Err(e) => {
turn.record_tool_error(e.to_string());
}
}
}
}
// If tool_auth returned awaiting_token, enter auth mode
// and short-circuit: return the instructions directly so
// the LLM doesn't get a chance to hallucinate tool calls.
if let Some((ext_name, instructions)) =
detect_auth_awaiting(&tc.name, &tool_result)
{
let auth_data = parse_auth_result(&tool_result);
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.enter_auth_mode(ext_name.clone());
}
}
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthRequired {
extension_name: ext_name,
instructions: Some(instructions.clone()),
auth_url: auth_data.auth_url,
setup_url: auth_data.setup_url,
},
&message.metadata,
)
.await;
return Ok(AgenticLoopResult::Response(instructions));
}
// Add tool result to context for next LLM call
let result_content = match tool_result {
Ok(output) => {
// Sanitize output before showing to LLM
let sanitized =
self.safety().sanitize_tool_output(&tc.name, &output);
self.safety().wrap_for_llm(
&tc.name,
&sanitized.content,
sanitized.was_modified,
)
}
Err(e) => format!("Error: {}", e),
};
context_messages.push(ChatMessage::tool_result(
&tc.id,
&tc.name,
result_content,
));
}
}
}
}
}
/// Execute a tool for chat (without full job context).
pub(super) async fn execute_chat_tool(
&self,
tool_name: &str,
params: &serde_json::Value,
job_ctx: &JobContext,
) -> Result<String, Error> {
let tool =
self.tools()
.get(tool_name)
.await
.ok_or_else(|| crate::error::ToolError::NotFound {
name: tool_name.to_string(),
})?;
// Validate tool parameters
let validation = self.safety().validator().validate_tool_params(params);
if !validation.is_valid {
let details = validation
.errors
.iter()
.map(|e| format!("{}: {}", e.field, e.message))
.collect::<Vec<_>>()
.join("; ");
return Err(crate::error::ToolError::InvalidParameters {
name: tool_name.to_string(),
reason: format!("Invalid tool parameters: {}", details),
}
.into());
}
tracing::debug!(
tool = %tool_name,
params = %params,
"Tool call started"
);
// Execute with per-tool timeout
let timeout = tool.execution_timeout();
let start = std::time::Instant::now();
let result = tokio::time::timeout(timeout, async {
tool.execute(params.clone(), job_ctx).await
})
.await;
let elapsed = start.elapsed();
match &result {
Ok(Ok(output)) => {
let result_str = serde_json::to_string(&output.result)
.unwrap_or_else(|_| "<serialize error>".to_string());
tracing::debug!(
tool = %tool_name,
elapsed_ms = elapsed.as_millis() as u64,
result = %result_str,
"Tool call succeeded"
);
}
Ok(Err(e)) => {
tracing::debug!(
tool = %tool_name,
elapsed_ms = elapsed.as_millis() as u64,
error = %e,
"Tool call failed"
);
}
Err(_) => {
tracing::debug!(
tool = %tool_name,
elapsed_ms = elapsed.as_millis() as u64,
timeout_secs = timeout.as_secs(),
"Tool call timed out"
);
}
}
let result = result
.map_err(|_| crate::error::ToolError::Timeout {
name: tool_name.to_string(),
timeout,
})?
.map_err(|e| crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: e.to_string(),
})?;
// Convert result to string
serde_json::to_string_pretty(&result.result).map_err(|e| {
crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: format!("Failed to serialize result: {}", e),
}
.into()
})
}
}
/// Parsed auth result fields for emitting StatusUpdate::AuthRequired.
pub(super) struct ParsedAuthData {
pub(super) auth_url: Option<String>,
pub(super) setup_url: Option<String>,
}
/// Extract auth_url and setup_url from a tool_auth result JSON string.
pub(super) fn parse_auth_result(result: &Result<String, Error>) -> ParsedAuthData {
let parsed = result
.as_ref()
.ok()
.and_then(|s| serde_json::from_str::<serde_json::Value>(s).ok());
ParsedAuthData {
auth_url: parsed
.as_ref()
.and_then(|v| v.get("auth_url"))
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
setup_url: parsed
.as_ref()
.and_then(|v| v.get("setup_url"))
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
}
}
/// Check if a tool_auth result indicates the extension is awaiting a token.
///
/// Returns `Some((extension_name, instructions))` if the tool result contains
/// `awaiting_token: true`, meaning the thread should enter auth mode.
pub(super) fn detect_auth_awaiting(
tool_name: &str,
result: &Result<String, Error>,
) -> Option<(String, String)> {
if tool_name != "tool_auth" && tool_name != "tool_activate" {
return None;
}
let output = result.as_ref().ok()?;
let parsed: serde_json::Value = serde_json::from_str(output).ok()?;
if parsed.get("awaiting_token") != Some(&serde_json::Value::Bool(true)) {
return None;
}
let name = parsed.get("name")?.as_str()?.to_string();
let instructions = parsed
.get("instructions")
.and_then(|v| v.as_str())
.unwrap_or("Please provide your API token/key.")
.to_string();
Some((name, instructions))
}
#[cfg(test)]
mod tests {
use crate::error::Error;
use super::detect_auth_awaiting;
#[test]
fn test_detect_auth_awaiting_positive() {
let result: Result<String, Error> = Ok(serde_json::json!({
"name": "telegram",
"kind": "WasmTool",
"awaiting_token": true,
"status": "awaiting_token",
"instructions": "Please provide your Telegram Bot API token."
})
.to_string());
let detected = detect_auth_awaiting("tool_auth", &result);
assert!(detected.is_some());
let (name, instructions) = detected.unwrap();
assert_eq!(name, "telegram");
assert!(instructions.contains("Telegram Bot API"));
}
#[test]
fn test_detect_auth_awaiting_not_awaiting() {
let result: Result<String, Error> = Ok(serde_json::json!({
"name": "telegram",
"kind": "WasmTool",
"awaiting_token": false,
"status": "authenticated"
})
.to_string());
assert!(detect_auth_awaiting("tool_auth", &result).is_none());
}
#[test]
fn test_detect_auth_awaiting_wrong_tool() {
let result: Result<String, Error> = Ok(serde_json::json!({
"name": "telegram",
"awaiting_token": true,
})
.to_string());
assert!(detect_auth_awaiting("tool_list", &result).is_none());
}
#[test]
fn test_detect_auth_awaiting_error_result() {
let result: Result<String, Error> =
Err(crate::error::ToolError::NotFound { name: "x".into() }.into());
assert!(detect_auth_awaiting("tool_auth", &result).is_none());
}
#[test]
fn test_detect_auth_awaiting_default_instructions() {
let result: Result<String, Error> = Ok(serde_json::json!({
"name": "custom_tool",
"awaiting_token": true,
"status": "awaiting_token"
})
.to_string());
let (_, instructions) = detect_auth_awaiting("tool_auth", &result).unwrap();
assert_eq!(instructions, "Please provide your API token/key.");
}
#[test]
fn test_detect_auth_awaiting_tool_activate() {
let result: Result<String, Error> = Ok(serde_json::json!({
"name": "slack",
"kind": "McpServer",
"awaiting_token": true,
"status": "awaiting_token",
"instructions": "Provide your Slack Bot token."
})
.to_string());
let detected = detect_auth_awaiting("tool_activate", &result);
assert!(detected.is_some());
let (name, instructions) = detected.unwrap();
assert_eq!(name, "slack");
assert!(instructions.contains("Slack Bot"));
}
#[test]
fn test_detect_auth_awaiting_tool_activate_not_awaiting() {
let result: Result<String, Error> = Ok(serde_json::json!({
"name": "slack",
"tools_loaded": ["slack_post_message"],
"message": "Activated"
})
.to_string());
assert!(detect_auth_awaiting("tool_activate", &result).is_none());
}
}
+4
View File
@@ -11,8 +11,11 @@
//! - Context compaction for long conversations
mod agent_loop;
mod commands;
pub mod compaction;
pub mod context_monitor;
pub mod cost_guard;
mod dispatcher;
mod heartbeat;
mod router;
pub mod routine;
@@ -23,6 +26,7 @@ pub mod session;
mod session_manager;
pub mod submission;
pub mod task;
mod thread_ops;
pub mod undo;
pub mod worker;
File diff suppressed because it is too large Load Diff
+20
View File
@@ -23,6 +23,10 @@ pub struct BootInfo {
pub claude_code_enabled: bool,
pub routines_enabled: bool,
pub channels: Vec<String>,
/// Public URL from a managed tunnel (e.g., "https://abc.ngrok.io").
pub tunnel_url: Option<String>,
/// Provider name for the managed tunnel (e.g., "ngrok").
pub tunnel_provider: Option<String>,
}
/// Print the boot screen to stdout.
@@ -116,6 +120,16 @@ pub fn print_boot_screen(info: &BootInfo) {
println!(" {dim}gateway{reset} {yellow_underline}{url}{reset}");
}
// Tunnel URL
if let Some(ref url) = info.tunnel_url {
let provider_tag = info
.tunnel_provider
.as_deref()
.map(|p| format!(" {dim}({p}){reset}"))
.unwrap_or_default();
println!(" {dim}tunnel{reset} {yellow_underline}{url}{reset}{provider_tag}");
}
println!();
println!("{border}");
println!();
@@ -151,6 +165,8 @@ mod tests {
"gateway".to_string(),
"telegram".to_string(),
],
tunnel_url: Some("https://abc123.ngrok.io".to_string()),
tunnel_provider: Some("ngrok".to_string()),
};
// Should not panic
print_boot_screen(&info);
@@ -176,6 +192,8 @@ mod tests {
claude_code_enabled: false,
routines_enabled: false,
channels: vec![],
tunnel_url: None,
tunnel_provider: None,
};
// Should not panic
print_boot_screen(&info);
@@ -201,6 +219,8 @@ mod tests {
claude_code_enabled: false,
routines_enabled: false,
channels: vec!["repl".to_string()],
tunnel_url: None,
tunnel_provider: None,
};
// Should not panic
print_boot_screen(&info);
+89
View File
@@ -11,6 +11,17 @@ use crate::settings::Settings;
#[derive(Subcommand, Debug, Clone)]
pub enum ConfigCommand {
/// Generate a default config.toml file
Init {
/// Output path (default: ~/.ironclaw/config.toml)
#[arg(short, long)]
output: Option<std::path::PathBuf>,
/// Overwrite existing file
#[arg(long)]
force: bool,
},
/// List all settings and their current values
List {
/// Show only settings matching this prefix (e.g., "agent", "heartbeat")
@@ -62,6 +73,7 @@ pub async fn run_config_command(cmd: ConfigCommand) -> anyhow::Result<()> {
let db_ref = db.as_deref();
match cmd {
ConfigCommand::Init { output, force } => init_toml(db_ref, output, force).await,
ConfigCommand::List { filter } => list_settings(db_ref, filter).await,
ConfigCommand::Get { path } => get_setting(db_ref, &path).await,
ConfigCommand::Set { path, value } => set_setting(db_ref, &path, &value).await,
@@ -188,6 +200,36 @@ async fn reset_setting(store: Option<&dyn crate::db::Database>, path: &str) -> a
Ok(())
}
/// Generate a default TOML config file.
async fn init_toml(
store: Option<&dyn crate::db::Database>,
output: Option<std::path::PathBuf>,
force: bool,
) -> anyhow::Result<()> {
let path = output.unwrap_or_else(Settings::default_toml_path);
if path.exists() && !force {
anyhow::bail!(
"Config file already exists: {}\nUse --force to overwrite.",
path.display()
);
}
// Start from current settings (DB or defaults) so the generated file
// reflects the user's existing configuration.
let settings = load_settings(store).await;
settings
.save_toml(&path)
.map_err(|e| anyhow::anyhow!("{}", e))?;
println!("Config file written to {}", path.display());
println!();
println!("Edit the file to customize settings.");
println!("Priority: env var > config.toml > database > defaults");
Ok(())
}
/// Show the settings storage info.
fn show_path(has_db: bool) -> anyhow::Result<()> {
if has_db {
@@ -200,6 +242,18 @@ fn show_path(has_db: bool) -> anyhow::Result<()> {
crate::bootstrap::ironclaw_env_path().display()
);
let toml_path = Settings::default_toml_path();
let toml_status = if toml_path.exists() {
"found"
} else {
"not found (run `ironclaw config init` to create)"
};
println!(
"TOML config: {} ({})",
toml_path.display(),
toml_status
);
Ok(())
}
@@ -230,4 +284,39 @@ mod tests {
settings.reset("agent.name").unwrap();
assert_eq!(settings.agent.name, "ironclaw");
}
#[tokio::test]
async fn init_toml_creates_file() {
let dir = tempdir().unwrap();
let path = dir.path().join("config.toml");
init_toml(None, Some(path.clone()), false).await.unwrap();
assert!(path.exists());
let content = std::fs::read_to_string(&path).unwrap();
assert!(content.contains("[agent]"));
}
#[tokio::test]
async fn init_toml_refuses_overwrite_without_force() {
let dir = tempdir().unwrap();
let path = dir.path().join("config.toml");
std::fs::write(&path, "existing").unwrap();
let result = init_toml(None, Some(path.clone()), false).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("already exists"));
}
#[tokio::test]
async fn init_toml_force_overwrites() {
let dir = tempdir().unwrap();
let path = dir.path().join("config.toml");
std::fs::write(&path, "old content").unwrap();
init_toml(None, Some(path.clone()), true).await.unwrap();
let content = std::fs::read_to_string(&path).unwrap();
assert!(content.contains("[agent]"));
}
}
+287
View File
@@ -0,0 +1,287 @@
//! `ironclaw doctor` - active health diagnostics.
//!
//! Probes external dependencies and validates configuration to surface
//! problems before they bite during normal operation. Each check reports
//! pass/fail with actionable guidance on failures.
use std::path::PathBuf;
/// Run all diagnostic checks and print results.
pub async fn run_doctor_command() -> anyhow::Result<()> {
println!("IronClaw Doctor");
println!("===============\n");
let mut passed = 0u32;
let mut failed = 0u32;
// ── Configuration checks ──────────────────────────────────
check(
"NEAR AI session",
check_nearai_session().await,
&mut passed,
&mut failed,
);
check(
"Database backend",
check_database().await,
&mut passed,
&mut failed,
);
check(
"Workspace directory",
check_workspace_dir(),
&mut passed,
&mut failed,
);
// ── External binary checks ────────────────────────────────
check(
"Docker",
check_binary("docker", &["--version"]),
&mut passed,
&mut failed,
);
check(
"cloudflared",
check_binary("cloudflared", &["--version"]),
&mut passed,
&mut failed,
);
check(
"ngrok",
check_binary("ngrok", &["version"]),
&mut passed,
&mut failed,
);
check(
"tailscale",
check_binary("tailscale", &["version"]),
&mut passed,
&mut failed,
);
// ── Summary ───────────────────────────────────────────────
println!();
println!(" {passed} passed, {failed} failed");
if failed > 0 {
println!("\n Some checks failed. This is normal if you don't use those features.");
}
Ok(())
}
// ── Individual checks ───────────────────────────────────────
fn check(name: &str, result: CheckResult, passed: &mut u32, failed: &mut u32) {
match result {
CheckResult::Pass(detail) => {
*passed += 1;
println!(" [pass] {name}: {detail}");
}
CheckResult::Fail(detail) => {
*failed += 1;
println!(" [FAIL] {name}: {detail}");
}
CheckResult::Skip(reason) => {
println!(" [skip] {name}: {reason}");
}
}
}
enum CheckResult {
Pass(String),
Fail(String),
Skip(String),
}
async fn check_nearai_session() -> CheckResult {
// Check if session file exists
let session_path = crate::llm::session::default_session_path();
if !session_path.exists() {
// Check for API key mode
if std::env::var("NEARAI_API_KEY").is_ok() {
return CheckResult::Pass("API key configured".into());
}
return CheckResult::Fail(format!(
"session file not found at {}. Run `ironclaw onboard`",
session_path.display()
));
}
// Verify the session file is readable and non-empty
match std::fs::read_to_string(&session_path) {
Ok(content) if content.trim().is_empty() => {
CheckResult::Fail("session file is empty".into())
}
Ok(_) => CheckResult::Pass(format!("session found ({})", session_path.display())),
Err(e) => CheckResult::Fail(format!("cannot read session file: {e}")),
}
}
async fn check_database() -> CheckResult {
let backend = std::env::var("DATABASE_BACKEND")
.ok()
.unwrap_or_else(|| "postgres".into());
match backend.as_str() {
"libsql" | "turso" | "sqlite" => {
let path = std::env::var("LIBSQL_PATH")
.map(PathBuf::from)
.unwrap_or_else(|_| crate::config::default_libsql_path());
if path.exists() {
CheckResult::Pass(format!("libSQL database exists ({})", path.display()))
} else {
CheckResult::Pass(format!(
"libSQL database not found at {} (will be created on first run)",
path.display()
))
}
}
_ => {
if std::env::var("DATABASE_URL").is_ok() {
// Try to connect
match try_pg_connect().await {
Ok(()) => CheckResult::Pass("PostgreSQL connected".into()),
Err(e) => CheckResult::Fail(format!("PostgreSQL connection failed: {e}")),
}
} else {
CheckResult::Fail("DATABASE_URL not set".into())
}
}
}
}
#[cfg(feature = "postgres")]
async fn try_pg_connect() -> Result<(), String> {
let url = std::env::var("DATABASE_URL").map_err(|_| "DATABASE_URL not set".to_string())?;
let config = deadpool_postgres::Config {
url: Some(url),
..Default::default()
};
let pool = config
.create_pool(
Some(deadpool_postgres::Runtime::Tokio1),
tokio_postgres::NoTls,
)
.map_err(|e| format!("pool error: {e}"))?;
let client = tokio::time::timeout(std::time::Duration::from_secs(5), pool.get())
.await
.map_err(|_| "connection timeout (5s)".to_string())?
.map_err(|e| format!("{e}"))?;
client
.execute("SELECT 1", &[])
.await
.map_err(|e| format!("{e}"))?;
Ok(())
}
#[cfg(not(feature = "postgres"))]
async fn try_pg_connect() -> Result<(), String> {
Err("postgres feature not compiled in".into())
}
fn check_workspace_dir() -> CheckResult {
let dir = dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw");
if dir.exists() {
if dir.is_dir() {
CheckResult::Pass(format!("{}", dir.display()))
} else {
CheckResult::Fail(format!("{} exists but is not a directory", dir.display()))
}
} else {
CheckResult::Pass(format!("{} will be created on first run", dir.display()))
}
}
fn check_binary(name: &str, args: &[&str]) -> CheckResult {
match std::process::Command::new(name)
.args(args)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.output()
{
Ok(output) => {
let version = String::from_utf8_lossy(&output.stdout);
let version = version.trim();
// Some tools print version to stderr
let version = if version.is_empty() {
let stderr = String::from_utf8_lossy(&output.stderr);
stderr.trim().lines().next().unwrap_or("").to_string()
} else {
version.lines().next().unwrap_or("").to_string()
};
if output.status.success() {
CheckResult::Pass(version)
} else {
CheckResult::Fail(format!("exited with {}", output.status))
}
}
Err(_) => CheckResult::Skip(format!("{name} not found in PATH")),
}
}
#[cfg(test)]
mod tests {
use crate::cli::doctor::*;
#[test]
fn check_binary_finds_sh() {
match check_binary("sh", &["-c", "echo ok"]) {
CheckResult::Pass(_) => {}
other => panic!("expected Pass for sh, got: {}", format_result(&other)),
}
}
#[test]
fn check_binary_skips_nonexistent() {
match check_binary("__ironclaw_nonexistent_binary__", &["--version"]) {
CheckResult::Skip(_) => {}
other => panic!(
"expected Skip for nonexistent binary, got: {}",
format_result(&other)
),
}
}
#[test]
fn check_workspace_dir_does_not_panic() {
let result = check_workspace_dir();
match result {
CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {}
}
}
#[tokio::test]
async fn check_nearai_session_does_not_panic() {
let result = check_nearai_session().await;
match result {
CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {}
}
}
fn format_result(r: &CheckResult) -> String {
match r {
CheckResult::Pass(s) => format!("Pass({s})"),
CheckResult::Fail(s) => format!("Fail({s})"),
CheckResult::Skip(s) => format!("Skip({s})"),
}
}
}
+13
View File
@@ -7,23 +7,29 @@
//! - Managing WASM tools (`tool install`, `tool list`, `tool remove`)
//! - Managing MCP servers (`mcp add`, `mcp auth`, `mcp list`, `mcp test`)
//! - Querying workspace memory (`memory search`, `memory read`, `memory write`)
//! - Managing OS service (`service install`, `service start`, `service stop`)
//! - Active health diagnostics (`doctor`)
//! - Checking system health (`status`)
mod config;
mod doctor;
mod mcp;
pub mod memory;
pub mod oauth_defaults;
mod pairing;
mod service;
pub mod status;
mod tool;
pub use config::{ConfigCommand, run_config_command};
pub use doctor::run_doctor_command;
pub use mcp::{McpCommand, run_mcp_command};
pub use memory::MemoryCommand;
#[cfg(feature = "postgres")]
pub use memory::run_memory_command;
pub use memory::run_memory_command_with_db;
pub use pairing::{PairingCommand, run_pairing_command, run_pairing_command_with_store};
pub use service::{ServiceCommand, run_service_command};
pub use status::run_status_command;
pub use tool::{ToolCommand, run_tool_command};
@@ -96,6 +102,13 @@ pub enum Command {
#[command(subcommand)]
Pairing(PairingCommand),
/// Manage OS service (launchd / systemd)
#[command(subcommand)]
Service(ServiceCommand),
/// Probe external dependencies and validate configuration
Doctor,
/// Show system health and diagnostics
Status,
+37
View File
@@ -0,0 +1,37 @@
//! CLI subcommand definitions for `ironclaw service`.
use clap::Subcommand;
use crate::service::ServiceAction;
#[derive(Subcommand, Debug, Clone)]
pub enum ServiceCommand {
/// Install the OS service (launchd on macOS, systemd on Linux).
Install,
/// Start the installed service.
Start,
/// Stop the running service.
Stop,
/// Show service status.
Status,
/// Uninstall the OS service and remove the unit file.
Uninstall,
}
impl ServiceCommand {
/// Convert the CLI variant into the domain action.
pub fn to_action(&self) -> ServiceAction {
match self {
ServiceCommand::Install => ServiceAction::Install,
ServiceCommand::Start => ServiceAction::Start,
ServiceCommand::Stop => ServiceAction::Stop,
ServiceCommand::Status => ServiceAction::Status,
ServiceCommand::Uninstall => ServiceAction::Uninstall,
}
}
}
/// Run the service command.
pub fn run_service_command(cmd: &ServiceCommand) -> anyhow::Result<()> {
crate::service::handle_command(&cmd.to_action())
}
+179 -6
View File
@@ -39,22 +39,32 @@ pub struct Config {
pub routines: RoutineConfig,
pub sandbox: SandboxModeConfig,
pub claude_code: ClaudeCodeConfig,
pub observability: crate::observability::ObservabilityConfig,
}
impl Config {
/// Load configuration from environment variables and the database.
///
/// Priority: env var > DB settings > default.
/// Priority: env var > TOML config file > DB settings > default.
/// This is the primary way to load config after DB is connected.
pub async fn from_db(
store: &dyn crate::db::Database,
user_id: &str,
) -> Result<Self, ConfigError> {
Self::from_db_with_toml(store, user_id, None).await
}
/// Load from DB with an optional TOML config file overlay.
pub async fn from_db_with_toml(
store: &dyn crate::db::Database,
user_id: &str,
toml_path: Option<&std::path::Path>,
) -> Result<Self, ConfigError> {
let _ = dotenvy::dotenv();
crate::bootstrap::load_ironclaw_env();
// Load all settings from DB into a Settings struct
let db_settings = match store.get_all_settings(user_id).await {
let mut db_settings = match store.get_all_settings(user_id).await {
Ok(map) => Settings::from_db_map(&map),
Err(e) => {
tracing::warn!("Failed to load settings from DB, using defaults: {}", e);
@@ -62,6 +72,9 @@ impl Config {
}
};
// Overlay TOML config file (values win over DB settings)
Self::apply_toml_overlay(&mut db_settings, toml_path)?;
Self::build(&db_settings).await
}
@@ -74,12 +87,63 @@ impl Config {
/// Loads both `./.env` (standard, higher priority) and `~/.ironclaw/.env`
/// (lower priority) via dotenvy, which never overwrites existing vars.
pub async fn from_env() -> Result<Self, ConfigError> {
Self::from_env_with_toml(None).await
}
/// Load from env with an optional TOML config file overlay.
pub async fn from_env_with_toml(
toml_path: Option<&std::path::Path>,
) -> Result<Self, ConfigError> {
let _ = dotenvy::dotenv();
crate::bootstrap::load_ironclaw_env();
let settings = Settings::load();
let mut settings = Settings::load();
// Overlay TOML config file (values win over JSON settings)
Self::apply_toml_overlay(&mut settings, toml_path)?;
Self::build(&settings).await
}
/// Load and merge a TOML config file into settings.
///
/// If `explicit_path` is `Some`, loads from that path (errors are fatal).
/// If `None`, tries the default path `~/.ironclaw/config.toml` (missing
/// file is silently ignored).
fn apply_toml_overlay(
settings: &mut Settings,
explicit_path: Option<&std::path::Path>,
) -> Result<(), ConfigError> {
let path = explicit_path
.map(std::path::PathBuf::from)
.unwrap_or_else(Settings::default_toml_path);
match Settings::load_toml(&path) {
Ok(Some(toml_settings)) => {
settings.merge_from(&toml_settings);
tracing::debug!("Loaded TOML config from {}", path.display());
}
Ok(None) => {
if explicit_path.is_some() {
return Err(ConfigError::ParseError(format!(
"Config file not found: {}",
path.display()
)));
}
}
Err(e) => {
if explicit_path.is_some() {
return Err(ConfigError::ParseError(format!(
"Failed to load config file {}: {}",
path.display(),
e
)));
}
tracing::warn!("Failed to load default config file: {}", e);
}
}
Ok(())
}
/// Build config from settings (shared by from_env and from_db).
async fn build(settings: &Settings) -> Result<Self, ConfigError> {
Ok(Self {
@@ -97,6 +161,9 @@ impl Config {
routines: RoutineConfig::resolve()?,
sandbox: SandboxModeConfig::resolve()?,
claude_code: ClaudeCodeConfig::resolve()?,
observability: crate::observability::ObservabilityConfig {
backend: std::env::var("OBSERVABILITY_BACKEND").unwrap_or_else(|_| "none".into()),
},
})
}
}
@@ -105,10 +172,21 @@ impl Config {
///
/// Used by channels and tools that need public webhook endpoints.
/// The tunnel URL is shared across all channels (Telegram, Slack, etc.).
///
/// Two modes:
/// - **Static URL** (`TUNNEL_URL`): set the public URL directly (manual tunnel)
/// - **Managed provider** (`TUNNEL_PROVIDER`): lifecycle-managed tunnel process
///
/// When a managed provider is configured _and_ no static URL is set,
/// the gateway starts the tunnel on boot and populates `public_url`.
#[derive(Debug, Clone, Default)]
pub struct TunnelConfig {
/// Public URL from tunnel provider (e.g., "https://abc123.ngrok.io").
/// Set statically via `TUNNEL_URL` or populated at runtime by a managed tunnel.
pub public_url: Option<String>,
/// Provider configuration for lifecycle-managed tunnels.
/// `None` when using a static URL or no tunnel at all.
pub provider: Option<crate::tunnel::TunnelProviderConfig>,
}
impl TunnelConfig {
@@ -125,12 +203,65 @@ impl TunnelConfig {
});
}
Ok(Self { public_url })
// Resolve managed tunnel provider config.
// Priority: env var > settings > default (none).
let provider_name = optional_env("TUNNEL_PROVIDER")?
.or_else(|| settings.tunnel.provider.clone())
.unwrap_or_default();
let provider = if provider_name.is_empty() || provider_name == "none" {
None
} else {
Some(crate::tunnel::TunnelProviderConfig {
provider: provider_name.clone(),
cloudflare: optional_env("TUNNEL_CF_TOKEN")?
.or_else(|| settings.tunnel.cf_token.clone())
.map(|token| crate::tunnel::CloudflareTunnelConfig { token }),
tailscale: Some(crate::tunnel::TailscaleTunnelConfig {
funnel: optional_env("TUNNEL_TS_FUNNEL")
.ok()
.flatten()
.map(|s| s == "true" || s == "1")
.unwrap_or(settings.tunnel.ts_funnel),
hostname: optional_env("TUNNEL_TS_HOSTNAME")
.ok()
.flatten()
.or_else(|| settings.tunnel.ts_hostname.clone()),
}),
ngrok: optional_env("TUNNEL_NGROK_TOKEN")?
.or_else(|| settings.tunnel.ngrok_token.clone())
.map(|auth_token| crate::tunnel::NgrokTunnelConfig {
auth_token,
domain: optional_env("TUNNEL_NGROK_DOMAIN")
.ok()
.flatten()
.or_else(|| settings.tunnel.ngrok_domain.clone()),
}),
custom: optional_env("TUNNEL_CUSTOM_COMMAND")?
.or_else(|| settings.tunnel.custom_command.clone())
.map(|start_command| crate::tunnel::CustomTunnelConfig {
start_command,
health_url: optional_env("TUNNEL_CUSTOM_HEALTH_URL")
.ok()
.flatten()
.or_else(|| settings.tunnel.custom_health_url.clone()),
url_pattern: optional_env("TUNNEL_CUSTOM_URL_PATTERN")
.ok()
.flatten()
.or_else(|| settings.tunnel.custom_url_pattern.clone()),
}),
})
};
Ok(Self {
public_url,
provider,
})
}
/// Check if a tunnel is configured.
/// Check if a tunnel is configured (static URL or managed provider).
pub fn is_enabled(&self) -> bool {
self.public_url.is_some()
self.public_url.is_some() || self.provider.is_some()
}
/// Get the webhook URL for a given path.
@@ -419,6 +550,19 @@ pub struct NearAiConfig {
/// With the default of 3, the provider makes up to 4 total attempts
/// (1 initial + 3 retries) before giving up.
pub max_retries: u32,
/// Consecutive transient failures before the circuit breaker opens.
/// None = disabled (default). E.g. 5 means after 5 consecutive failures
/// all requests are rejected until recovery timeout elapses.
pub circuit_breaker_threshold: Option<u32>,
/// How long (seconds) the circuit stays open before allowing a probe (default: 30).
pub circuit_breaker_recovery_secs: u64,
/// Enable in-memory response caching for `complete()` calls.
/// Saves tokens on repeated prompts within a session. Default: false.
pub response_cache_enabled: bool,
/// TTL in seconds for cached responses (default: 3600 = 1 hour).
pub response_cache_ttl_secs: u64,
/// Max cached responses before LRU eviction (default: 1000).
pub response_cache_max_entries: usize,
/// Cooldown duration in seconds for the failover provider (default: 300).
/// When a provider accumulates enough consecutive failures it is skipped
/// for this many seconds.
@@ -485,6 +629,17 @@ impl LlmConfig {
api_key: nearai_api_key,
fallback_model: optional_env("NEARAI_FALLBACK_MODEL")?,
max_retries: parse_optional_env("NEARAI_MAX_RETRIES", 3)?,
circuit_breaker_threshold: optional_env("CIRCUIT_BREAKER_THRESHOLD")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "CIRCUIT_BREAKER_THRESHOLD".to_string(),
message: format!("must be a positive integer: {e}"),
})?,
circuit_breaker_recovery_secs: parse_optional_env("CIRCUIT_BREAKER_RECOVERY_SECS", 30)?,
response_cache_enabled: parse_optional_env("RESPONSE_CACHE_ENABLED", false)?,
response_cache_ttl_secs: parse_optional_env("RESPONSE_CACHE_TTL_SECS", 3600)?,
response_cache_max_entries: parse_optional_env("RESPONSE_CACHE_MAX_ENTRIES", 1000)?,
failover_cooldown_secs: parse_optional_env("LLM_FAILOVER_COOLDOWN_SECS", 300)?,
failover_cooldown_threshold: parse_optional_env("LLM_FAILOVER_THRESHOLD", 3)?,
};
@@ -755,6 +910,10 @@ pub struct AgentConfig {
pub session_idle_timeout: Duration,
/// Allow chat to use filesystem/shell tools directly (bypass sandbox).
pub allow_local_tools: bool,
/// Maximum daily LLM spend in cents (e.g. 10000 = $100). None = unlimited.
pub max_cost_per_day_cents: Option<u64>,
/// Maximum LLM/tool actions per hour. None = unlimited.
pub max_actions_per_hour: Option<u64>,
}
impl AgentConfig {
@@ -833,6 +992,20 @@ impl AgentConfig {
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or(false),
max_cost_per_day_cents: optional_env("MAX_COST_PER_DAY_CENTS")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "MAX_COST_PER_DAY_CENTS".to_string(),
message: format!("must be a positive integer: {e}"),
})?,
max_actions_per_hour: optional_env("MAX_ACTIONS_PER_HOUR")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "MAX_ACTIONS_PER_HOUR".to_string(),
message: format!("must be a positive integer: {e}"),
})?,
})
}
}
+3
View File
@@ -53,15 +53,18 @@ pub mod extensions;
pub mod history;
pub mod hooks;
pub mod llm;
pub mod observability;
pub mod orchestrator;
pub mod pairing;
pub mod safety;
pub mod sandbox;
pub mod secrets;
pub mod service;
pub mod settings;
pub mod setup;
pub mod tools;
pub mod tracing_fmt;
pub mod tunnel;
pub mod util;
pub mod worker;
pub mod workspace;
+674
View File
@@ -0,0 +1,674 @@
//! Circuit breaker for LLM providers.
//!
//! Wraps any `LlmProvider` with a state machine that trips open after
//! consecutive transient failures, preventing request storms against a
//! degraded backend. Automatically probes for recovery via half-open state.
//!
//! ```text
//! Closed ──(failures >= threshold)──► Open
//! ▲ │
//! │ (recovery timeout)
//! │ ▼
//! └──(probe succeeds)──── HalfOpen ──(probe fails)──► Open
//! ```
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use rust_decimal::Decimal;
use tokio::sync::Mutex;
use crate::error::LlmError;
use crate::llm::provider::{
CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest,
ToolCompletionResponse,
};
/// Configuration for the circuit breaker.
#[derive(Debug, Clone)]
pub struct CircuitBreakerConfig {
/// Consecutive transient failures before the circuit opens.
pub failure_threshold: u32,
/// How long the circuit stays open before allowing a probe.
pub recovery_timeout: Duration,
/// Successful probes needed in half-open to close the circuit.
pub half_open_successes_needed: u32,
}
impl Default for CircuitBreakerConfig {
fn default() -> Self {
Self {
failure_threshold: 5,
recovery_timeout: Duration::from_secs(30),
half_open_successes_needed: 2,
}
}
}
/// Circuit breaker states.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
/// Normal operation; tracking consecutive failures.
Closed,
/// Rejecting all calls; waiting for recovery timeout to elapse.
Open,
/// Allowing probe calls to test whether the backend recovered.
HalfOpen,
}
/// Internal mutable state.
struct BreakerState {
state: CircuitState,
consecutive_failures: u32,
opened_at: Option<Instant>,
half_open_successes: u32,
}
impl BreakerState {
fn new() -> Self {
Self {
state: CircuitState::Closed,
consecutive_failures: 0,
opened_at: None,
half_open_successes: 0,
}
}
}
/// Wraps an `LlmProvider` with circuit breaker protection.
///
/// Tracks consecutive transient failures. After `failure_threshold` failures
/// the circuit opens and all requests are rejected for `recovery_timeout`.
/// After that timeout a probe call is allowed through (half-open); if it
/// succeeds the circuit closes, otherwise it reopens.
pub struct CircuitBreakerProvider {
inner: Arc<dyn LlmProvider>,
state: Mutex<BreakerState>,
config: CircuitBreakerConfig,
}
impl CircuitBreakerProvider {
pub fn new(inner: Arc<dyn LlmProvider>, config: CircuitBreakerConfig) -> Self {
Self {
inner,
state: Mutex::new(BreakerState::new()),
config,
}
}
/// Current circuit state (for observability / health checks).
pub async fn circuit_state(&self) -> CircuitState {
self.state.lock().await.state
}
/// Number of consecutive failures recorded so far.
pub async fn consecutive_failures(&self) -> u32 {
self.state.lock().await.consecutive_failures
}
/// Pre-flight: is a call allowed right now?
async fn check_allowed(&self) -> Result<(), LlmError> {
let mut state = self.state.lock().await;
match state.state {
CircuitState::Closed | CircuitState::HalfOpen => Ok(()),
CircuitState::Open => {
if let Some(opened_at) = state.opened_at {
if opened_at.elapsed() >= self.config.recovery_timeout {
state.state = CircuitState::HalfOpen;
state.half_open_successes = 0;
tracing::info!(
provider = self.inner.model_name(),
"Circuit breaker: Open -> HalfOpen, allowing probe"
);
Ok(())
} else {
let remaining = self.config.recovery_timeout - opened_at.elapsed();
Err(LlmError::RequestFailed {
provider: self.inner.model_name().to_string(),
reason: format!(
"Circuit breaker open ({} consecutive failures, \
recovery in {:.0}s)",
state.consecutive_failures,
remaining.as_secs_f64()
),
})
}
} else {
// opened_at should always be Some when Open; recover gracefully
state.state = CircuitState::Closed;
Ok(())
}
}
}
}
/// Record a successful call.
async fn record_success(&self) {
let mut state = self.state.lock().await;
match state.state {
CircuitState::Closed => {
state.consecutive_failures = 0;
}
CircuitState::HalfOpen => {
state.half_open_successes += 1;
if state.half_open_successes >= self.config.half_open_successes_needed {
state.state = CircuitState::Closed;
state.consecutive_failures = 0;
state.opened_at = None;
tracing::info!(
provider = self.inner.model_name(),
"Circuit breaker: HalfOpen -> Closed (recovered)"
);
}
}
CircuitState::Open => {
// Shouldn't get here (check_allowed blocks Open), but recover
state.state = CircuitState::Closed;
state.consecutive_failures = 0;
state.opened_at = None;
}
}
}
/// Record a failed call; only transient errors count toward the threshold.
async fn record_failure(&self, err: &LlmError) {
if !is_transient(err) {
return;
}
let mut state = self.state.lock().await;
match state.state {
CircuitState::Closed => {
state.consecutive_failures += 1;
if state.consecutive_failures >= self.config.failure_threshold {
state.state = CircuitState::Open;
state.opened_at = Some(Instant::now());
tracing::warn!(
provider = self.inner.model_name(),
failures = state.consecutive_failures,
"Circuit breaker: Closed -> Open"
);
}
}
CircuitState::HalfOpen => {
state.state = CircuitState::Open;
state.opened_at = Some(Instant::now());
state.half_open_successes = 0;
tracing::warn!(
provider = self.inner.model_name(),
"Circuit breaker: HalfOpen -> Open (probe failed)"
);
}
CircuitState::Open => {}
}
}
}
/// Returns `true` for errors that indicate the provider is degraded
/// (server errors, rate limits, network failures, auth infrastructure down).
///
/// Client errors (wrong model, bad credentials, context overflow) are NOT
/// transient: they are the caller's problem, not a sign of backend trouble.
fn is_transient(err: &LlmError) -> bool {
matches!(
err,
LlmError::RequestFailed { .. }
| LlmError::RateLimited { .. }
| LlmError::InvalidResponse { .. }
| LlmError::SessionExpired { .. }
| LlmError::SessionRenewalFailed { .. }
| LlmError::Http(_)
| LlmError::Json(_)
| LlmError::Io(_)
)
}
#[async_trait]
impl LlmProvider for CircuitBreakerProvider {
fn model_name(&self) -> &str {
self.inner.model_name()
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
self.inner.cost_per_token()
}
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
self.check_allowed().await?;
match self.inner.complete(request).await {
Ok(resp) => {
self.record_success().await;
Ok(resp)
}
Err(err) => {
self.record_failure(&err).await;
Err(err)
}
}
}
async fn complete_with_tools(
&self,
request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
self.check_allowed().await?;
match self.inner.complete_with_tools(request).await {
Ok(resp) => {
self.record_success().await;
Ok(resp)
}
Err(err) => {
self.record_failure(&err).await;
Err(err)
}
}
}
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
self.inner.list_models().await
}
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
self.inner.model_metadata().await
}
fn active_model_name(&self) -> String {
self.inner.active_model_name()
}
fn set_model(&self, model: &str) -> Result<(), LlmError> {
self.inner.set_model(model)
}
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
self.inner.seed_response_chain(thread_id, response_id)
}
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
self.inner.get_response_chain_id(thread_id)
}
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
self.inner.calculate_cost(input_tokens, output_tokens)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, Ordering};
use crate::llm::provider::{CompletionResponse, FinishReason, ToolCompletionResponse};
/// A test stub that either always succeeds or always fails with a
/// configurable error. The `should_fail` flag can be flipped at
/// runtime for half-open recovery tests.
struct StubProvider {
name: String,
should_fail: AtomicBool,
error_kind: StubError,
}
#[derive(Clone)]
enum StubError {
Transient,
NonTransient,
}
impl StubProvider {
fn always_ok(name: &str) -> Arc<Self> {
Arc::new(Self {
name: name.to_string(),
should_fail: AtomicBool::new(false),
error_kind: StubError::Transient,
})
}
fn always_fail(name: &str) -> Arc<Self> {
Arc::new(Self {
name: name.to_string(),
should_fail: AtomicBool::new(true),
error_kind: StubError::Transient,
})
}
fn always_fail_non_transient(name: &str) -> Arc<Self> {
Arc::new(Self {
name: name.to_string(),
should_fail: AtomicBool::new(true),
error_kind: StubError::NonTransient,
})
}
fn set_failing(&self, fail: bool) {
self.should_fail.store(fail, Ordering::Relaxed);
}
fn make_error(&self) -> LlmError {
match self.error_kind {
StubError::Transient => LlmError::RequestFailed {
provider: self.name.clone(),
reason: "server error".to_string(),
},
StubError::NonTransient => LlmError::ContextLengthExceeded {
used: 100_000,
limit: 50_000,
},
}
}
fn ok_response() -> CompletionResponse {
CompletionResponse {
content: "ok".to_string(),
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
}
}
fn ok_tool_response() -> ToolCompletionResponse {
ToolCompletionResponse {
content: Some("ok".to_string()),
tool_calls: vec![],
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
}
}
}
#[async_trait]
impl LlmProvider for StubProvider {
fn model_name(&self) -> &str {
&self.name
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
(Decimal::ZERO, Decimal::ZERO)
}
async fn complete(
&self,
_request: CompletionRequest,
) -> Result<CompletionResponse, LlmError> {
if self.should_fail.load(Ordering::Relaxed) {
Err(self.make_error())
} else {
Ok(Self::ok_response())
}
}
async fn complete_with_tools(
&self,
_request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
if self.should_fail.load(Ordering::Relaxed) {
Err(self.make_error())
} else {
Ok(Self::ok_tool_response())
}
}
}
fn make_request() -> CompletionRequest {
CompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")])
}
fn make_tool_request() -> ToolCompletionRequest {
ToolCompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")], vec![])
}
fn fast_config(threshold: u32) -> CircuitBreakerConfig {
CircuitBreakerConfig {
failure_threshold: threshold,
recovery_timeout: Duration::from_millis(50),
half_open_successes_needed: 1,
}
}
// -- State machine tests --
#[tokio::test]
async fn closed_allows_calls_and_resets_on_success() {
let stub = StubProvider::always_ok("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(3));
let resp = cb.complete(make_request()).await;
assert!(resp.is_ok());
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
assert_eq!(cb.consecutive_failures().await, 0);
}
#[tokio::test]
async fn failures_accumulate_then_trip_to_open() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(3));
// First 2 failures: still closed
for i in 0..2 {
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
assert_eq!(cb.consecutive_failures().await, i + 1);
}
// 3rd failure: trips to open
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
}
#[tokio::test]
async fn open_rejects_immediately() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(
stub,
CircuitBreakerConfig {
failure_threshold: 1,
recovery_timeout: Duration::from_secs(60),
half_open_successes_needed: 1,
},
);
// Trip the breaker
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
// Next call should fail with circuit breaker message
let err = cb.complete(make_request()).await.unwrap_err();
match err {
LlmError::RequestFailed { reason, .. } => {
assert!(
reason.contains("Circuit breaker open"),
"Expected circuit breaker message, got: {}",
reason
);
}
other => panic!("Expected RequestFailed, got: {:?}", other),
}
}
#[tokio::test]
async fn recovery_timeout_transitions_to_half_open() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(1));
// Trip to open
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
// Wait for recovery timeout
tokio::time::sleep(Duration::from_millis(60)).await;
// Next call should transition to half-open (and fail, since stub fails)
let _ = cb.complete(make_request()).await;
// Failed probe sends it back to Open
assert_eq!(cb.circuit_state().await, CircuitState::Open);
}
#[tokio::test]
async fn half_open_success_closes_circuit() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub.clone(), fast_config(1));
// Trip to open
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
// Wait for recovery, then make the stub succeed
tokio::time::sleep(Duration::from_millis(60)).await;
stub.set_failing(false);
// Probe should succeed, closing the circuit
let resp = cb.complete(make_request()).await;
assert!(resp.is_ok());
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
assert_eq!(cb.consecutive_failures().await, 0);
}
#[tokio::test]
async fn half_open_failure_reopens_circuit() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(1));
// Trip to open
let _ = cb.complete(make_request()).await;
// Wait for recovery timeout
tokio::time::sleep(Duration::from_millis(60)).await;
// Probe fails (stub still failing)
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
}
#[tokio::test]
async fn non_transient_errors_do_not_trip_breaker() {
let stub = StubProvider::always_fail_non_transient("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(1));
// ContextLengthExceeded is not transient; breaker should stay closed
for _ in 0..5 {
let _ = cb.complete(make_request()).await;
}
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
assert_eq!(cb.consecutive_failures().await, 0);
}
#[tokio::test]
async fn success_resets_failure_count() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub.clone(), fast_config(3));
// Accumulate 2 failures
let _ = cb.complete(make_request()).await;
let _ = cb.complete(make_request()).await;
assert_eq!(cb.consecutive_failures().await, 2);
// One success resets the counter
stub.set_failing(false);
let resp = cb.complete(make_request()).await;
assert!(resp.is_ok());
assert_eq!(cb.consecutive_failures().await, 0);
}
#[tokio::test]
async fn complete_with_tools_uses_same_breaker_logic() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(2));
let _ = cb.complete_with_tools(make_tool_request()).await;
let _ = cb.complete_with_tools(make_tool_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
}
#[tokio::test]
async fn multiple_half_open_successes_needed() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(
stub.clone(),
CircuitBreakerConfig {
failure_threshold: 1,
recovery_timeout: Duration::from_millis(50),
half_open_successes_needed: 3,
},
);
// Trip to open
let _ = cb.complete(make_request()).await;
// Wait and flip to succeed
tokio::time::sleep(Duration::from_millis(60)).await;
stub.set_failing(false);
// First probe: half-open, success but not enough yet
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::HalfOpen);
// Second probe: still half-open
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::HalfOpen);
// Third probe: closes
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
}
// -- Error classification tests --
#[test]
fn transient_classification() {
// Transient
assert!(is_transient(&LlmError::RequestFailed {
provider: "p".into(),
reason: "err".into(),
}));
assert!(is_transient(&LlmError::RateLimited {
provider: "p".into(),
retry_after: None,
}));
assert!(is_transient(&LlmError::InvalidResponse {
provider: "p".into(),
reason: "bad".into(),
}));
assert!(is_transient(&LlmError::SessionExpired {
provider: "p".into(),
}));
assert!(is_transient(&LlmError::SessionRenewalFailed {
provider: "p".into(),
reason: "timeout".into(),
}));
assert!(is_transient(&LlmError::Io(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset"
))));
// NOT transient
assert!(!is_transient(&LlmError::AuthFailed {
provider: "p".into(),
}));
assert!(!is_transient(&LlmError::ContextLengthExceeded {
used: 100_000,
limit: 50_000,
}));
assert!(!is_transient(&LlmError::ModelNotAvailable {
provider: "p".into(),
model: "m".into(),
}));
}
// -- Passthrough delegation tests --
#[tokio::test]
async fn passthrough_methods_delegate_to_inner() {
let stub = StubProvider::always_ok("my-model");
let cb = CircuitBreakerProvider::new(stub, fast_config(3));
assert_eq!(cb.model_name(), "my-model");
assert_eq!(cb.active_model_name(), "my-model");
assert_eq!(cb.cost_per_token(), (Decimal::ZERO, Decimal::ZERO));
assert_eq!(cb.calculate_cost(100, 50), Decimal::ZERO);
}
}
+10 -1
View File
@@ -7,16 +7,19 @@
//! - **Ollama**: Local model inference
//! - **OpenAI-compatible**: Any endpoint that speaks the OpenAI API
mod costs;
pub mod circuit_breaker;
pub mod costs;
pub mod failover;
mod nearai;
mod nearai_chat;
mod provider;
mod reasoning;
pub mod response_cache;
mod retry;
mod rig_adapter;
pub mod session;
pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider};
pub use failover::{CooldownConfig, FailoverProvider};
pub use nearai::{ModelInfo, NearAiProvider};
pub use nearai_chat::NearAiChatProvider;
@@ -28,6 +31,7 @@ pub use reasoning::{
ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, TokenUsage,
ToolSelection,
};
pub use response_cache::{CachedProvider, ResponseCacheConfig};
pub use rig_adapter::RigAdapter;
pub use session::{SessionConfig, SessionManager, create_session_manager};
@@ -235,6 +239,11 @@ mod tests {
api_key: None,
fallback_model: None,
max_retries: 3,
circuit_breaker_threshold: None,
circuit_breaker_recovery_secs: 30,
response_cache_enabled: false,
response_cache_ttl_secs: 3600,
response_cache_max_entries: 1000,
failover_cooldown_secs: 300,
failover_cooldown_threshold: 3,
}
+526
View File
@@ -0,0 +1,526 @@
//! In-memory LLM response cache with TTL and LRU eviction.
//!
//! Wraps any [`LlmProvider`] and caches [`complete()`] responses keyed
//! by a SHA-256 hash of the messages and model name. Tool-calling
//! requests are never cached since they can trigger side effects.
//!
//! ```text
//! ┌──────────────────────────────────────────────────┐
//! │ CachedProvider │
//! │ complete() ──► cache lookup ──► hit? return │
//! │ miss? call inner │
//! │ store response │
//! │ │
//! │ complete_with_tools() ──► always call inner │
//! └──────────────────────────────────────────────────┘
//! ```
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use rust_decimal::Decimal;
use sha2::{Digest, Sha256};
use tokio::sync::Mutex;
use crate::error::LlmError;
use crate::llm::provider::{
CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest,
ToolCompletionResponse,
};
/// Configuration for the response cache.
#[derive(Debug, Clone)]
pub struct ResponseCacheConfig {
/// Time-to-live for cache entries.
pub ttl: Duration,
/// Maximum number of cached entries before LRU eviction.
pub max_entries: usize,
}
impl Default for ResponseCacheConfig {
fn default() -> Self {
Self {
ttl: Duration::from_secs(3600), // 1 hour
max_entries: 1000,
}
}
}
struct CacheEntry {
response: CompletionResponse,
created_at: Instant,
last_accessed: Instant,
hit_count: u64,
}
/// LLM provider wrapper that caches `complete()` responses.
///
/// Tool completion requests are always forwarded without caching since
/// tool calls can have side effects that should not be replayed.
pub struct CachedProvider {
inner: Arc<dyn LlmProvider>,
cache: Mutex<HashMap<String, CacheEntry>>,
config: ResponseCacheConfig,
}
impl CachedProvider {
/// Wrap an existing provider with response caching.
pub fn new(inner: Arc<dyn LlmProvider>, config: ResponseCacheConfig) -> Self {
Self {
inner,
cache: Mutex::new(HashMap::new()),
config,
}
}
/// Number of entries currently in the cache.
pub async fn len(&self) -> usize {
self.cache.lock().await.len()
}
/// Whether the cache is empty.
pub async fn is_empty(&self) -> bool {
self.cache.lock().await.is_empty()
}
/// Total cache hits across all entries.
pub async fn total_hits(&self) -> u64 {
self.cache.lock().await.values().map(|e| e.hit_count).sum()
}
/// Clear all cached entries.
pub async fn clear(&self) {
self.cache.lock().await.clear();
}
}
/// Build a deterministic cache key from a completion request.
///
/// Hashes the model name, messages, and response-affecting parameters
/// (max_tokens, temperature, stop_sequences) via SHA-256. Two requests
/// with identical content and parameters produce the same key.
fn cache_key(model: &str, request: &CompletionRequest) -> String {
let mut hasher = Sha256::new();
hasher.update(model.as_bytes());
hasher.update(b"|");
// Messages are Serialize, so we can deterministically hash them.
// serde_json produces stable output for the same input structure.
if let Ok(json) = serde_json::to_string(&request.messages) {
hasher.update(json.as_bytes());
}
// Include response-affecting parameters so different temperatures,
// max_tokens, or stop sequences produce distinct cache keys.
hasher.update(b"|");
if let Some(max_tokens) = request.max_tokens {
hasher.update(max_tokens.to_le_bytes());
}
hasher.update(b"|");
if let Some(temp) = request.temperature {
hasher.update(temp.to_le_bytes());
}
hasher.update(b"|");
if let Some(ref stops) = request.stop_sequences {
for s in stops {
hasher.update(s.as_bytes());
hasher.update(b"\x00");
}
}
format!("{:x}", hasher.finalize())
}
#[async_trait]
impl LlmProvider for CachedProvider {
fn model_name(&self) -> &str {
self.inner.model_name()
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
self.inner.cost_per_token()
}
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
let key = cache_key(self.inner.model_name(), &request);
let now = Instant::now();
// Check cache
{
let mut guard = self.cache.lock().await;
if let Some(entry) = guard.get_mut(&key) {
if now.duration_since(entry.created_at) < self.config.ttl {
entry.last_accessed = now;
entry.hit_count += 1;
tracing::debug!(hits = entry.hit_count, "response cache hit");
return Ok(entry.response.clone());
}
// Expired, remove it
guard.remove(&key);
}
}
// Cache miss, call the real provider
let response = self.inner.complete(request).await?;
// Store in cache
{
let mut guard = self.cache.lock().await;
// Evict expired entries
guard.retain(|_, entry| now.duration_since(entry.created_at) < self.config.ttl);
// LRU eviction if over capacity
while guard.len() >= self.config.max_entries {
let oldest_key = guard
.iter()
.min_by_key(|(_, entry)| entry.last_accessed)
.map(|(k, _)| k.clone());
if let Some(k) = oldest_key {
guard.remove(&k);
} else {
break;
}
}
guard.insert(
key,
CacheEntry {
response: response.clone(),
created_at: now,
last_accessed: now,
hit_count: 0,
},
);
}
Ok(response)
}
async fn complete_with_tools(
&self,
request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
// Never cache tool calls; they can trigger side effects.
self.inner.complete_with_tools(request).await
}
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
self.inner.list_models().await
}
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
self.inner.model_metadata().await
}
fn active_model_name(&self) -> String {
self.inner.active_model_name()
}
fn set_model(&self, model: &str) -> Result<(), LlmError> {
self.inner.set_model(model)
}
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
self.inner.seed_response_chain(thread_id, response_id);
}
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
self.inner.get_response_chain_id(thread_id)
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use crate::llm::provider::{ChatMessage, FinishReason};
use crate::llm::response_cache::*;
/// Controllable stub provider for testing cache behavior.
struct StubProvider {
call_count: AtomicU32,
should_fail: AtomicBool,
}
impl StubProvider {
fn new() -> Self {
Self {
call_count: AtomicU32::new(0),
should_fail: AtomicBool::new(false),
}
}
fn calls(&self) -> u32 {
self.call_count.load(Ordering::Relaxed)
}
}
#[async_trait]
impl LlmProvider for StubProvider {
fn model_name(&self) -> &str {
"stub-model"
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
(Decimal::ZERO, Decimal::ZERO)
}
async fn complete(
&self,
_request: CompletionRequest,
) -> Result<CompletionResponse, LlmError> {
self.call_count.fetch_add(1, Ordering::Relaxed);
if self.should_fail.load(Ordering::Relaxed) {
return Err(LlmError::RequestFailed {
provider: "stub".into(),
reason: "forced failure".into(),
});
}
Ok(CompletionResponse {
content: "cached response".into(),
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
})
}
async fn complete_with_tools(
&self,
_request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
self.call_count.fetch_add(1, Ordering::Relaxed);
Ok(ToolCompletionResponse {
content: Some("tool response".into()),
tool_calls: vec![],
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
})
}
}
fn simple_request() -> CompletionRequest {
CompletionRequest {
messages: vec![ChatMessage::user("hello")],
max_tokens: None,
temperature: None,
stop_sequences: None,
metadata: Default::default(),
}
}
fn different_request() -> CompletionRequest {
CompletionRequest {
messages: vec![ChatMessage::user("goodbye")],
max_tokens: None,
temperature: None,
stop_sequences: None,
metadata: Default::default(),
}
}
#[test]
fn cache_key_is_deterministic() {
let req = simple_request();
let k1 = cache_key("model-a", &req);
let k2 = cache_key("model-a", &req);
assert_eq!(k1, k2);
assert_eq!(k1.len(), 64); // SHA-256 hex
}
#[test]
fn cache_key_varies_by_model() {
let req = simple_request();
let k1 = cache_key("model-a", &req);
let k2 = cache_key("model-b", &req);
assert_ne!(k1, k2);
}
#[test]
fn cache_key_varies_by_messages() {
let k1 = cache_key("model-a", &simple_request());
let k2 = cache_key("model-a", &different_request());
assert_ne!(k1, k2);
}
#[test]
fn cache_key_varies_by_temperature() {
let mut req_a = simple_request();
req_a.temperature = Some(0.0);
let mut req_b = simple_request();
req_b.temperature = Some(1.0);
assert_ne!(cache_key("m", &req_a), cache_key("m", &req_b));
}
#[test]
fn cache_key_varies_by_max_tokens() {
let mut req_a = simple_request();
req_a.max_tokens = Some(100);
let mut req_b = simple_request();
req_b.max_tokens = Some(500);
assert_ne!(cache_key("m", &req_a), cache_key("m", &req_b));
}
#[tokio::test]
async fn cache_hit_avoids_provider_call() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(
stub.clone(),
ResponseCacheConfig {
ttl: Duration::from_secs(60),
max_entries: 100,
},
);
// First call: cache miss
let r1 = cached.complete(simple_request()).await.unwrap();
assert_eq!(stub.calls(), 1);
assert_eq!(r1.content, "cached response");
// Second call: cache hit
let r2 = cached.complete(simple_request()).await.unwrap();
assert_eq!(stub.calls(), 1); // still 1
assert_eq!(r2.content, "cached response");
assert_eq!(cached.total_hits().await, 1);
}
#[tokio::test]
async fn different_messages_get_different_entries() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default());
cached.complete(simple_request()).await.unwrap();
cached.complete(different_request()).await.unwrap();
assert_eq!(stub.calls(), 2);
assert_eq!(cached.len().await, 2);
}
#[tokio::test]
async fn expired_entries_are_evicted() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(
stub.clone(),
ResponseCacheConfig {
ttl: Duration::from_millis(1),
max_entries: 100,
},
);
cached.complete(simple_request()).await.unwrap();
assert_eq!(stub.calls(), 1);
// Wait for TTL to expire
tokio::time::sleep(Duration::from_millis(10)).await;
// Should be a cache miss now
cached.complete(simple_request()).await.unwrap();
assert_eq!(stub.calls(), 2);
}
#[tokio::test]
async fn lru_eviction_removes_oldest() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(
stub.clone(),
ResponseCacheConfig {
ttl: Duration::from_secs(60),
max_entries: 2,
},
);
// Fill cache with 2 entries
cached.complete(simple_request()).await.unwrap();
cached.complete(different_request()).await.unwrap();
assert_eq!(cached.len().await, 2);
// Add a third: should evict the oldest
let third = CompletionRequest {
messages: vec![ChatMessage::user("third")],
max_tokens: None,
temperature: None,
stop_sequences: None,
metadata: Default::default(),
};
cached.complete(third).await.unwrap();
assert_eq!(cached.len().await, 2);
assert_eq!(stub.calls(), 3);
}
#[tokio::test]
async fn tool_calls_are_never_cached() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default());
let req = ToolCompletionRequest {
messages: vec![ChatMessage::user("use tool")],
tools: vec![],
max_tokens: None,
temperature: None,
tool_choice: None,
metadata: Default::default(),
};
cached.complete_with_tools(req.clone()).await.unwrap();
cached.complete_with_tools(req).await.unwrap();
// Both should have called through
assert_eq!(stub.calls(), 2);
assert!(cached.is_empty().await);
}
#[tokio::test]
async fn provider_errors_are_not_cached() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(
stub.clone(),
ResponseCacheConfig {
ttl: Duration::from_secs(60),
max_entries: 100,
},
);
stub.should_fail.store(true, Ordering::Relaxed);
let result = cached.complete(simple_request()).await;
assert!(result.is_err());
assert!(cached.is_empty().await);
// After fixing the provider, should succeed and cache
stub.should_fail.store(false, Ordering::Relaxed);
cached.complete(simple_request()).await.unwrap();
assert_eq!(cached.len().await, 1);
}
#[tokio::test]
async fn clear_empties_cache() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default());
cached.complete(simple_request()).await.unwrap();
assert_eq!(cached.len().await, 1);
cached.clear().await;
assert!(cached.is_empty().await);
}
#[test]
fn default_config_is_reasonable() {
let cfg = ResponseCacheConfig::default();
assert_eq!(cfg.ttl, Duration::from_secs(3600));
assert_eq!(cfg.max_entries, 1000);
}
#[tokio::test]
async fn delegates_model_name() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default());
assert_eq!(cached.model_name(), "stub-model");
}
}
+3 -3
View File
@@ -432,9 +432,9 @@ impl SessionManager {
.get_setting(&user_id, "nearai.session_token")
.await
.map_err(|e| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!("DB query failed: {}", e),
})? {
provider: "nearai".to_string(),
reason: format!("DB query failed: {}", e),
})? {
value
} else {
tracing::warn!(
+141 -7
View File
@@ -17,15 +17,18 @@ use ironclaw::{
web::log_layer::{LogBroadcaster, WebLogLayer},
},
cli::{
Cli, Command, run_mcp_command, run_pairing_command, run_status_command, run_tool_command,
Cli, Command, run_mcp_command, run_pairing_command, run_service_command,
run_status_command, run_tool_command,
},
config::Config,
context::ContextManager,
extensions::ExtensionManager,
hooks::HookRegistry,
llm::{
CooldownConfig, FailoverProvider, LlmProvider, SessionConfig, create_cheap_llm_provider,
create_llm_provider, create_llm_provider_with_config, create_session_manager,
CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig,
FailoverProvider, LlmProvider, ResponseCacheConfig, SessionConfig,
create_cheap_llm_provider, create_llm_provider, create_llm_provider_with_config,
create_session_manager,
},
orchestrator::{
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
@@ -152,6 +155,24 @@ async fn main() -> anyhow::Result<()> {
return run_pairing_command(pairing_cmd.clone()).map_err(|e| anyhow::anyhow!("{}", e));
}
Some(Command::Service(service_cmd)) => {
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
)
.init();
return run_service_command(service_cmd);
}
Some(Command::Doctor) => {
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
)
.init();
return ironclaw::cli::run_doctor_command().await;
}
Some(Command::Status) => {
tracing_subscriber::fmt()
.with_env_filter(
@@ -286,8 +307,9 @@ async fn main() -> anyhow::Result<()> {
wizard.run().await?;
}
// Load initial config from env + disk (before DB is available)
let mut config = match Config::from_env().await {
// Load initial config from env + disk + optional TOML (before DB is available)
let toml_path = cli.config.as_deref();
let mut config = match Config::from_env_with_toml(toml_path).await {
Ok(c) => c,
Err(ironclaw::error::ConfigError::MissingRequired { key, hint }) => {
eprintln!("Configuration error: Missing required setting '{}'", key);
@@ -429,7 +451,7 @@ async fn main() -> anyhow::Result<()> {
}
// Reload config from DB now that we have a connection.
match Config::from_db(db.as_ref(), "default").await {
match Config::from_db_with_toml(db.as_ref(), "default", toml_path).await {
Ok(db_config) => {
config = db_config;
tracing::info!("Configuration reloaded from database");
@@ -504,7 +526,7 @@ async fn main() -> anyhow::Result<()> {
// Re-resolve LlmConfig now that secrets overlay has been populated
if let Some(ref db_ref) = db {
match Config::from_db(db_ref.as_ref(), "default").await {
match Config::from_db_with_toml(db_ref.as_ref(), "default", toml_path).await {
Ok(refreshed) => {
config = refreshed;
tracing::debug!("LlmConfig re-resolved after secret injection");
@@ -516,6 +538,62 @@ async fn main() -> anyhow::Result<()> {
}
}
// Start managed tunnel if configured and no static URL is already set.
//
// The tunnel process runs in the background, exposing the local gateway
// port to the internet. The resulting public URL is injected into
// config.tunnel.public_url so channels and extensions pick it up.
let active_tunnel: Option<Box<dyn ironclaw::tunnel::Tunnel>> =
if config.tunnel.public_url.is_some() {
tracing::info!(
"Static tunnel URL in use: {}",
config.tunnel.public_url.as_deref().unwrap_or("?")
);
None
} else if let Some(ref provider_config) = config.tunnel.provider {
let gateway_port = config
.channels
.gateway
.as_ref()
.map(|g| g.port)
.unwrap_or(3000);
let gateway_host = config
.channels
.gateway
.as_ref()
.map(|g| g.host.as_str())
.unwrap_or("127.0.0.1");
match ironclaw::tunnel::create_tunnel(provider_config) {
Ok(Some(tunnel)) => {
tracing::info!(
"Starting {} tunnel on {}:{}...",
tunnel.name(),
gateway_host,
gateway_port
);
match tunnel.start(gateway_host, gateway_port).await {
Ok(url) => {
tracing::info!("Tunnel started: {}", url);
config.tunnel.public_url = Some(url);
Some(tunnel)
}
Err(e) => {
tracing::error!("Failed to start tunnel: {}", e);
None
}
}
}
Ok(None) => None,
Err(e) => {
tracing::error!("Failed to create tunnel: {}", e);
None
}
}
} else {
None
};
// Initialize LLM provider (clone session so we can reuse it for embeddings)
let llm = create_llm_provider(&config.llm, session.clone())?;
tracing::info!("LLM provider initialized: {}", llm.model_name());
@@ -550,6 +628,42 @@ async fn main() -> anyhow::Result<()> {
llm
};
// Wrap in circuit breaker if configured
let llm: Arc<dyn LlmProvider> =
if let Some(threshold) = config.llm.nearai.circuit_breaker_threshold {
let cb_config = CircuitBreakerConfig {
failure_threshold: threshold,
recovery_timeout: std::time::Duration::from_secs(
config.llm.nearai.circuit_breaker_recovery_secs,
),
..CircuitBreakerConfig::default()
};
tracing::info!(
threshold,
recovery_secs = config.llm.nearai.circuit_breaker_recovery_secs,
"LLM circuit breaker enabled"
);
Arc::new(CircuitBreakerProvider::new(llm, cb_config))
} else {
llm
};
// Wrap in response cache if configured
let llm: Arc<dyn LlmProvider> = if config.llm.nearai.response_cache_enabled {
let rc_config = ResponseCacheConfig {
ttl: std::time::Duration::from_secs(config.llm.nearai.response_cache_ttl_secs),
max_entries: config.llm.nearai.response_cache_max_entries,
};
tracing::info!(
ttl_secs = config.llm.nearai.response_cache_ttl_secs,
max_entries = config.llm.nearai.response_cache_max_entries,
"LLM response cache enabled"
);
Arc::new(CachedProvider::new(llm, rc_config))
} else {
llm
};
// Initialize cheap LLM provider for lightweight tasks (heartbeat, evaluation)
let cheap_llm = create_cheap_llm_provider(&config.llm, session.clone())?;
if let Some(ref cheap) = cheap_llm {
@@ -1220,6 +1334,12 @@ async fn main() -> anyhow::Result<()> {
let boot_cheap_model = cheap_llm.as_ref().map(|c| c.model_name().to_string());
// Create and run the agent
let cost_guard = Arc::new(ironclaw::agent::cost_guard::CostGuard::new(
ironclaw::agent::cost_guard::CostGuardConfig {
max_cost_per_day_cents: config.agent.max_cost_per_day_cents,
max_actions_per_hour: config.agent.max_actions_per_hour,
},
));
let deps = AgentDeps {
store: db,
llm,
@@ -1229,6 +1349,7 @@ async fn main() -> anyhow::Result<()> {
workspace,
extension_manager,
hooks,
cost_guard,
};
let agent = Agent::new(
config.agent.clone(),
@@ -1270,6 +1391,11 @@ async fn main() -> anyhow::Result<()> {
claude_code_enabled: config.claude_code.enabled,
routines_enabled: config.routines.enabled,
channels: channel_names,
tunnel_url: active_tunnel
.as_ref()
.and_then(|t| t.public_url())
.or_else(|| config.tunnel.public_url.clone()),
tunnel_provider: active_tunnel.as_ref().map(|t| t.name().to_string()),
};
ironclaw::boot_screen::print_boot_screen(&boot_info);
}
@@ -1282,6 +1408,14 @@ async fn main() -> anyhow::Result<()> {
server.shutdown().await;
}
// Stop managed tunnel if one was started
if let Some(tunnel) = active_tunnel {
tracing::info!("Stopping {} tunnel...", tunnel.name());
if let Err(e) = tunnel.stop().await {
tracing::warn!("Failed to stop tunnel cleanly: {}", e);
}
}
tracing::info!("Agent shutdown complete");
Ok(())
}
+181
View File
@@ -0,0 +1,181 @@
//! Tracing-based observer that emits structured log events.
//!
//! Uses the existing `tracing` infrastructure so events appear alongside
//! normal application logs, with no extra dependencies. Good for local
//! development and debugging.
use crate::observability::traits::{Observer, ObserverEvent, ObserverMetric};
/// Observer that logs events and metrics via `tracing`.
pub struct LogObserver;
impl Observer for LogObserver {
fn record_event(&self, event: &ObserverEvent) {
match event {
ObserverEvent::AgentStart { provider, model } => {
tracing::info!(provider, model, "observer: agent.start");
}
ObserverEvent::LlmRequest {
provider,
model,
message_count,
} => {
tracing::info!(provider, model, message_count, "observer: llm.request");
}
ObserverEvent::LlmResponse {
provider,
model,
duration,
success,
error_message,
} => {
tracing::info!(
provider,
model,
duration_ms = duration.as_millis() as u64,
success,
error = error_message.as_deref().unwrap_or(""),
"observer: llm.response"
);
}
ObserverEvent::ToolCallStart { tool } => {
tracing::info!(tool, "observer: tool.start");
}
ObserverEvent::ToolCallEnd {
tool,
duration,
success,
} => {
tracing::info!(
tool,
duration_ms = duration.as_millis() as u64,
success,
"observer: tool.end"
);
}
ObserverEvent::TurnComplete => {
tracing::info!("observer: turn.complete");
}
ObserverEvent::ChannelMessage { channel, direction } => {
tracing::info!(channel, direction, "observer: channel.message");
}
ObserverEvent::HeartbeatTick => {
tracing::debug!("observer: heartbeat.tick");
}
ObserverEvent::AgentEnd {
duration,
tokens_used,
} => {
tracing::info!(
duration_secs = duration.as_secs_f64(),
tokens_used = tokens_used.unwrap_or(0),
"observer: agent.end"
);
}
ObserverEvent::Error { component, message } => {
tracing::warn!(component, error = message.as_str(), "observer: error");
}
}
}
fn record_metric(&self, metric: &ObserverMetric) {
match metric {
ObserverMetric::RequestLatency(d) => {
tracing::debug!(
latency_ms = d.as_millis() as u64,
"observer: metric.request_latency"
);
}
ObserverMetric::TokensUsed(n) => {
tracing::debug!(tokens = n, "observer: metric.tokens_used");
}
ObserverMetric::ActiveJobs(n) => {
tracing::debug!(active_jobs = n, "observer: metric.active_jobs");
}
ObserverMetric::QueueDepth(n) => {
tracing::debug!(queue_depth = n, "observer: metric.queue_depth");
}
}
}
fn name(&self) -> &str {
"log"
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use crate::observability::log::LogObserver;
use crate::observability::traits::*;
#[test]
fn name_is_log() {
assert_eq!(LogObserver.name(), "log");
}
#[test]
fn record_event_does_not_panic() {
let obs = LogObserver;
obs.record_event(&ObserverEvent::AgentStart {
provider: "nearai".into(),
model: "test".into(),
});
obs.record_event(&ObserverEvent::LlmRequest {
provider: "nearai".into(),
model: "test".into(),
message_count: 5,
});
obs.record_event(&ObserverEvent::LlmResponse {
provider: "nearai".into(),
model: "test".into(),
duration: Duration::from_millis(150),
success: true,
error_message: None,
});
obs.record_event(&ObserverEvent::LlmResponse {
provider: "nearai".into(),
model: "test".into(),
duration: Duration::from_millis(1500),
success: false,
error_message: Some("timeout".into()),
});
obs.record_event(&ObserverEvent::ToolCallStart {
tool: "shell".into(),
});
obs.record_event(&ObserverEvent::ToolCallEnd {
tool: "shell".into(),
duration: Duration::from_millis(20),
success: true,
});
obs.record_event(&ObserverEvent::TurnComplete);
obs.record_event(&ObserverEvent::ChannelMessage {
channel: "tui".into(),
direction: "inbound".into(),
});
obs.record_event(&ObserverEvent::HeartbeatTick);
obs.record_event(&ObserverEvent::AgentEnd {
duration: Duration::from_secs(30),
tokens_used: Some(2500),
});
obs.record_event(&ObserverEvent::Error {
component: "llm".into(),
message: "connection refused".into(),
});
}
#[test]
fn record_metric_does_not_panic() {
let obs = LogObserver;
obs.record_metric(&ObserverMetric::RequestLatency(Duration::from_millis(200)));
obs.record_metric(&ObserverMetric::TokensUsed(1000));
obs.record_metric(&ObserverMetric::ActiveJobs(5));
obs.record_metric(&ObserverMetric::QueueDepth(12));
}
#[test]
fn flush_does_not_panic() {
LogObserver.flush();
}
}
+105
View File
@@ -0,0 +1,105 @@
//! Observability subsystem: trait-based event and metric recording.
//!
//! Provides a pluggable [`Observer`] trait with multiple backends:
//!
//! | Backend | Description |
//! |---------|-------------|
//! | `noop` | Zero overhead, discards everything (default) |
//! | `log` | Emits structured events via `tracing` |
//! | `multi` | Fan-out to multiple backends simultaneously |
//!
//! The [`create_observer`] factory builds the right backend from
//! [`ObservabilityConfig`]. Future backends (OpenTelemetry, Prometheus)
//! can be added by implementing [`Observer`].
mod log;
mod multi;
mod noop;
pub mod traits;
pub use self::log::LogObserver;
pub use self::multi::MultiObserver;
pub use self::noop::NoopObserver;
pub use self::traits::{Observer, ObserverEvent, ObserverMetric};
/// Configuration for the observability backend.
#[derive(Debug, Clone)]
pub struct ObservabilityConfig {
/// Backend name: "none", "noop", "log".
pub backend: String,
}
impl Default for ObservabilityConfig {
fn default() -> Self {
Self {
backend: "none".into(),
}
}
}
/// Create an observer from configuration.
///
/// Returns a [`NoopObserver`] for "none"/"noop" (or unknown values),
/// and a [`LogObserver`] for "log".
pub fn create_observer(config: &ObservabilityConfig) -> Box<dyn Observer> {
match config.backend.as_str() {
"log" => Box::new(LogObserver),
_ => Box::new(NoopObserver),
}
}
#[cfg(test)]
mod tests {
use crate::observability::*;
#[test]
fn default_config_is_none() {
let cfg = ObservabilityConfig::default();
assert_eq!(cfg.backend, "none");
}
#[test]
fn factory_returns_noop_for_none() {
let cfg = ObservabilityConfig {
backend: "none".into(),
};
let obs = create_observer(&cfg);
assert_eq!(obs.name(), "noop");
}
#[test]
fn factory_returns_noop_for_empty() {
let cfg = ObservabilityConfig {
backend: String::new(),
};
let obs = create_observer(&cfg);
assert_eq!(obs.name(), "noop");
}
#[test]
fn factory_returns_noop_for_unknown() {
let cfg = ObservabilityConfig {
backend: "prometheus".into(),
};
let obs = create_observer(&cfg);
assert_eq!(obs.name(), "noop");
}
#[test]
fn factory_returns_log_for_log() {
let cfg = ObservabilityConfig {
backend: "log".into(),
};
let obs = create_observer(&cfg);
assert_eq!(obs.name(), "log");
}
#[test]
fn factory_returns_noop_for_noop() {
let cfg = ObservabilityConfig {
backend: "noop".into(),
};
let obs = create_observer(&cfg);
assert_eq!(obs.name(), "noop");
}
}
+139
View File
@@ -0,0 +1,139 @@
//! Fan-out observer that dispatches to multiple backends.
//!
//! Useful for combining backends, e.g. log + OpenTelemetry simultaneously.
use crate::observability::traits::{Observer, ObserverEvent, ObserverMetric};
/// Dispatches events and metrics to all inner observers.
pub struct MultiObserver {
observers: Vec<Box<dyn Observer>>,
}
impl MultiObserver {
/// Create from a list of observers. If the list is empty the result
/// behaves like a noop.
pub fn new(observers: Vec<Box<dyn Observer>>) -> Self {
Self { observers }
}
}
impl Observer for MultiObserver {
fn record_event(&self, event: &ObserverEvent) {
for obs in &self.observers {
obs.record_event(event);
}
}
fn record_metric(&self, metric: &ObserverMetric) {
for obs in &self.observers {
obs.record_metric(metric);
}
}
fn flush(&self) {
for obs in &self.observers {
obs.flush();
}
}
fn name(&self) -> &str {
"multi"
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use crate::observability::multi::MultiObserver;
use crate::observability::traits::*;
/// Test observer that counts calls via shared atomic counters.
struct CountingObserver {
events: Arc<AtomicUsize>,
metrics: Arc<AtomicUsize>,
flushes: Arc<AtomicUsize>,
}
impl CountingObserver {
fn new() -> (Self, Arc<AtomicUsize>, Arc<AtomicUsize>, Arc<AtomicUsize>) {
let events = Arc::new(AtomicUsize::new(0));
let metrics = Arc::new(AtomicUsize::new(0));
let flushes = Arc::new(AtomicUsize::new(0));
(
Self {
events: Arc::clone(&events),
metrics: Arc::clone(&metrics),
flushes: Arc::clone(&flushes),
},
events,
metrics,
flushes,
)
}
}
impl Observer for CountingObserver {
fn record_event(&self, _event: &ObserverEvent) {
self.events.fetch_add(1, Ordering::Relaxed);
}
fn record_metric(&self, _metric: &ObserverMetric) {
self.metrics.fetch_add(1, Ordering::Relaxed);
}
fn flush(&self) {
self.flushes.fetch_add(1, Ordering::Relaxed);
}
fn name(&self) -> &str {
"counting"
}
}
#[test]
fn name_is_multi() {
let multi = MultiObserver::new(vec![]);
assert_eq!(multi.name(), "multi");
}
#[test]
fn empty_multi_does_not_panic() {
let multi = MultiObserver::new(vec![]);
multi.record_event(&ObserverEvent::TurnComplete);
multi.record_metric(&ObserverMetric::TokensUsed(100));
multi.flush();
}
#[test]
fn dispatches_to_all_observers() {
let (a, a_events, a_metrics, a_flushes) = CountingObserver::new();
let (b, b_events, b_metrics, b_flushes) = CountingObserver::new();
let multi = MultiObserver::new(vec![Box::new(a), Box::new(b)]);
multi.record_event(&ObserverEvent::TurnComplete);
multi.record_event(&ObserverEvent::HeartbeatTick);
multi.record_metric(&ObserverMetric::TokensUsed(50));
multi.flush();
assert_eq!(a_events.load(Ordering::Relaxed), 2);
assert_eq!(a_metrics.load(Ordering::Relaxed), 1);
assert_eq!(a_flushes.load(Ordering::Relaxed), 1);
assert_eq!(b_events.load(Ordering::Relaxed), 2);
assert_eq!(b_metrics.load(Ordering::Relaxed), 1);
assert_eq!(b_flushes.load(Ordering::Relaxed), 1);
}
#[test]
fn single_observer_works() {
let (obs, events, _, _) = CountingObserver::new();
let multi = MultiObserver::new(vec![Box::new(obs)]);
multi.record_event(&ObserverEvent::AgentEnd {
duration: Duration::from_secs(1),
tokens_used: None,
});
assert_eq!(events.load(Ordering::Relaxed), 1);
}
}
+60
View File
@@ -0,0 +1,60 @@
//! Zero-overhead no-op observer.
//!
//! Default backend when observability is disabled. All methods compile to
//! nothing, so there is zero runtime cost.
use crate::observability::traits::{Observer, ObserverEvent, ObserverMetric};
/// Observer that discards all events and metrics.
pub struct NoopObserver;
impl Observer for NoopObserver {
#[inline(always)]
fn record_event(&self, _event: &ObserverEvent) {}
#[inline(always)]
fn record_metric(&self, _metric: &ObserverMetric) {}
fn name(&self) -> &str {
"noop"
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use crate::observability::traits::*;
use crate::observability::noop::NoopObserver;
#[test]
fn name_is_noop() {
assert_eq!(NoopObserver.name(), "noop");
}
#[test]
fn record_event_does_not_panic() {
let obs = NoopObserver;
obs.record_event(&ObserverEvent::TurnComplete);
obs.record_event(&ObserverEvent::HeartbeatTick);
obs.record_event(&ObserverEvent::AgentStart {
provider: "x".into(),
model: "y".into(),
});
}
#[test]
fn record_metric_does_not_panic() {
let obs = NoopObserver;
obs.record_metric(&ObserverMetric::TokensUsed(100));
obs.record_metric(&ObserverMetric::RequestLatency(Duration::from_millis(50)));
obs.record_metric(&ObserverMetric::ActiveJobs(2));
obs.record_metric(&ObserverMetric::QueueDepth(0));
}
#[test]
fn flush_does_not_panic() {
NoopObserver.flush();
}
}
+143
View File
@@ -0,0 +1,143 @@
//! Core observer trait and event/metric types.
use std::time::Duration;
/// Provider-agnostic observer for agent lifecycle events and metrics.
///
/// Implementations can log to tracing, export to OpenTelemetry, write to
/// Prometheus, or do nothing at all. The agent records events at key
/// lifecycle points and the observer decides what to do with them.
///
/// Thread-safe and cheaply cloneable behind `Arc<dyn Observer>`.
pub trait Observer: Send + Sync {
/// Record a discrete lifecycle event.
fn record_event(&self, event: &ObserverEvent);
/// Record a numeric metric sample.
fn record_metric(&self, metric: &ObserverMetric);
/// Flush any buffered data (e.g. OTLP batch exporter). No-op by default.
fn flush(&self) {}
/// Human-readable backend name (e.g. "noop", "log", "otel").
fn name(&self) -> &str;
}
/// Discrete lifecycle events the agent can emit.
#[derive(Debug, Clone)]
pub enum ObserverEvent {
/// Agent started processing.
AgentStart { provider: String, model: String },
/// An LLM request was sent.
LlmRequest {
provider: String,
model: String,
message_count: usize,
},
/// An LLM response was received.
LlmResponse {
provider: String,
model: String,
duration: Duration,
success: bool,
error_message: Option<String>,
},
/// A tool call is about to start.
ToolCallStart { tool: String },
/// A tool call finished.
ToolCallEnd {
tool: String,
duration: Duration,
success: bool,
},
/// One reasoning turn completed.
TurnComplete,
/// A message was sent or received on a channel.
ChannelMessage { channel: String, direction: String },
/// The heartbeat system ran a tick.
HeartbeatTick,
/// Agent finished processing.
AgentEnd {
duration: Duration,
tokens_used: Option<u64>,
},
/// An error occurred in a component.
Error { component: String, message: String },
}
/// Numeric metric samples.
#[derive(Debug, Clone)]
pub enum ObserverMetric {
/// Latency of a single request (histogram-style).
RequestLatency(Duration),
/// Cumulative tokens consumed.
TokensUsed(u64),
/// Current number of active jobs (gauge).
ActiveJobs(u64),
/// Current message queue depth (gauge).
QueueDepth(u64),
}
#[cfg(test)]
mod tests {
use crate::observability::traits::*;
#[test]
fn event_variants_are_constructible() {
let _ = ObserverEvent::AgentStart {
provider: "nearai".into(),
model: "test".into(),
};
let _ = ObserverEvent::LlmRequest {
provider: "nearai".into(),
model: "test".into(),
message_count: 3,
};
let _ = ObserverEvent::LlmResponse {
provider: "nearai".into(),
model: "test".into(),
duration: Duration::from_millis(100),
success: true,
error_message: None,
};
let _ = ObserverEvent::ToolCallStart {
tool: "echo".into(),
};
let _ = ObserverEvent::ToolCallEnd {
tool: "echo".into(),
duration: Duration::from_millis(5),
success: true,
};
let _ = ObserverEvent::TurnComplete;
let _ = ObserverEvent::ChannelMessage {
channel: "tui".into(),
direction: "inbound".into(),
};
let _ = ObserverEvent::HeartbeatTick;
let _ = ObserverEvent::AgentEnd {
duration: Duration::from_secs(10),
tokens_used: Some(1500),
};
let _ = ObserverEvent::Error {
component: "llm".into(),
message: "timeout".into(),
};
}
#[test]
fn metric_variants_are_constructible() {
let _ = ObserverMetric::RequestLatency(Duration::from_millis(200));
let _ = ObserverMetric::TokensUsed(500);
let _ = ObserverMetric::ActiveJobs(3);
let _ = ObserverMetric::QueueDepth(10);
}
}
+357
View File
@@ -0,0 +1,357 @@
//! OS service management for running IronClaw as a daemon.
//!
//! Generates and manages platform-native service definitions:
//! - **macOS**: launchd plist at `~/Library/LaunchAgents/com.ironclaw.daemon.plist`
//! - **Linux**: systemd user unit at `~/.config/systemd/user/ironclaw.service`
//!
//! The installed service runs `ironclaw run` (the default agent mode) and is
//! configured to restart automatically on failure.
use std::path::PathBuf;
use std::process::Command;
use anyhow::{Context, Result, bail};
const SERVICE_LABEL: &str = "com.ironclaw.daemon";
const SYSTEMD_UNIT: &str = "ironclaw.service";
// ── Public dispatch ─────────────────────────────────────────────
/// Route a service subcommand to the appropriate handler.
pub fn handle_command(command: &ServiceAction) -> Result<()> {
match command {
ServiceAction::Install => install(),
ServiceAction::Start => start(),
ServiceAction::Stop => stop(),
ServiceAction::Status => status(),
ServiceAction::Uninstall => uninstall(),
}
}
/// The five service lifecycle actions.
#[derive(Debug, Clone)]
pub enum ServiceAction {
Install,
Start,
Stop,
Status,
Uninstall,
}
// ── Install ─────────────────────────────────────────────────────
fn install() -> Result<()> {
if cfg!(target_os = "macos") {
install_macos()
} else if cfg!(target_os = "linux") {
install_linux()
} else {
bail!("Service management is only supported on macOS and Linux");
}
}
fn install_macos() -> Result<()> {
let file = macos_plist_path()?;
if let Some(parent) = file.parent() {
std::fs::create_dir_all(parent)?;
}
let exe = std::env::current_exe().context("failed to resolve current executable")?;
let logs_dir = ironclaw_logs_dir()?;
std::fs::create_dir_all(&logs_dir)?;
let stdout = logs_dir.join("daemon.stdout.log");
let stderr = logs_dir.join("daemon.stderr.log");
let plist = format!(
r#"<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>Label</key>
<string>{label}</string>
<key>ProgramArguments</key>
<array>
<string>{exe}</string>
<string>run</string>
</array>
<key>RunAtLoad</key>
<true/>
<key>KeepAlive</key>
<true/>
<key>StandardOutPath</key>
<string>{stdout}</string>
<key>StandardErrorPath</key>
<string>{stderr}</string>
</dict>
</plist>
"#,
label = SERVICE_LABEL,
exe = xml_escape(&exe.display().to_string()),
stdout = xml_escape(&stdout.display().to_string()),
stderr = xml_escape(&stderr.display().to_string()),
);
std::fs::write(&file, plist)?;
println!("Installed launchd service: {}", file.display());
println!(" Start with: ironclaw service start");
Ok(())
}
fn install_linux() -> Result<()> {
let file = linux_unit_path()?;
if let Some(parent) = file.parent() {
std::fs::create_dir_all(parent)?;
}
let exe = std::env::current_exe().context("failed to resolve current executable")?;
let unit = format!(
"[Unit]\n\
Description=IronClaw daemon\n\
After=network.target\n\
\n\
[Service]\n\
Type=simple\n\
ExecStart=\"{exe}\" run\n\
Restart=always\n\
RestartSec=3\n\
\n\
[Install]\n\
WantedBy=default.target\n",
exe = exe.display(),
);
std::fs::write(&file, unit)?;
run_checked(Command::new("systemctl").args(["--user", "daemon-reload"])).ok();
run_checked(Command::new("systemctl").args(["--user", "enable", SYSTEMD_UNIT])).ok();
println!("Installed systemd user service: {}", file.display());
println!(" Start with: ironclaw service start");
Ok(())
}
// ── Start ───────────────────────────────────────────────────────
fn start() -> Result<()> {
if cfg!(target_os = "macos") {
let plist = macos_plist_path()?;
if !plist.exists() {
bail!("Service not installed. Run `ironclaw service install` first.");
}
run_checked(Command::new("launchctl").arg("load").arg("-w").arg(&plist))?;
run_checked(Command::new("launchctl").arg("start").arg(SERVICE_LABEL))?;
println!("Service started");
Ok(())
} else if cfg!(target_os = "linux") {
run_checked(Command::new("systemctl").args(["--user", "daemon-reload"]))?;
run_checked(Command::new("systemctl").args(["--user", "start", SYSTEMD_UNIT]))?;
println!("Service started");
Ok(())
} else {
bail!("Service management is only supported on macOS and Linux");
}
}
// ── Stop ────────────────────────────────────────────────────────
fn stop() -> Result<()> {
if cfg!(target_os = "macos") {
let plist = macos_plist_path()?;
run_checked(Command::new("launchctl").arg("stop").arg(SERVICE_LABEL)).ok();
run_checked(
Command::new("launchctl")
.arg("unload")
.arg("-w")
.arg(&plist),
)
.ok();
println!("Service stopped");
Ok(())
} else if cfg!(target_os = "linux") {
run_checked(Command::new("systemctl").args(["--user", "stop", SYSTEMD_UNIT])).ok();
println!("Service stopped");
Ok(())
} else {
bail!("Service management is only supported on macOS and Linux");
}
}
// ── Status ──────────────────────────────────────────────────────
fn status() -> Result<()> {
if cfg!(target_os = "macos") {
let out = run_capture(Command::new("launchctl").arg("list"))?;
let running = out.lines().any(|line| line.contains(SERVICE_LABEL));
println!(
"Service: {}",
if running {
"running/loaded"
} else {
"not loaded"
}
);
println!("Unit: {}", macos_plist_path()?.display());
Ok(())
} else if cfg!(target_os = "linux") {
let state =
run_capture(Command::new("systemctl").args(["--user", "is-active", SYSTEMD_UNIT]))
.unwrap_or_else(|_| "unknown".into());
println!("Service state: {}", state.trim());
println!("Unit: {}", linux_unit_path()?.display());
Ok(())
} else {
bail!("Service management is only supported on macOS and Linux");
}
}
// ── Uninstall ───────────────────────────────────────────────────
fn uninstall() -> Result<()> {
// Stop first (ignore errors, service might not be running)
stop().ok();
if cfg!(target_os = "macos") {
let file = macos_plist_path()?;
if file.exists() {
std::fs::remove_file(&file)
.with_context(|| format!("failed to remove {}", file.display()))?;
}
println!("Service uninstalled ({})", file.display());
Ok(())
} else if cfg!(target_os = "linux") {
let file = linux_unit_path()?;
if file.exists() {
std::fs::remove_file(&file)
.with_context(|| format!("failed to remove {}", file.display()))?;
}
run_checked(Command::new("systemctl").args(["--user", "daemon-reload"])).ok();
println!("Service uninstalled ({})", file.display());
Ok(())
} else {
bail!("Service management is only supported on macOS and Linux");
}
}
// ── Path helpers ────────────────────────────────────────────────
fn macos_plist_path() -> Result<PathBuf> {
let home = dirs::home_dir().context("could not find home directory")?;
Ok(home
.join("Library")
.join("LaunchAgents")
.join(format!("{SERVICE_LABEL}.plist")))
}
fn linux_unit_path() -> Result<PathBuf> {
let home = dirs::home_dir().context("could not find home directory")?;
Ok(home
.join(".config")
.join("systemd")
.join("user")
.join(SYSTEMD_UNIT))
}
fn ironclaw_logs_dir() -> Result<PathBuf> {
let home = dirs::home_dir().context("could not find home directory")?;
Ok(home.join(".ironclaw").join("logs"))
}
// ── Shell helpers ───────────────────────────────────────────────
fn run_checked(command: &mut Command) -> Result<()> {
let output = command.output().context("failed to spawn command")?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
bail!("command failed: {}", stderr.trim());
}
Ok(())
}
fn run_capture(command: &mut Command) -> Result<String> {
let output = command.output().context("failed to spawn command")?;
let mut text = String::from_utf8_lossy(&output.stdout).to_string();
if text.trim().is_empty() {
text = String::from_utf8_lossy(&output.stderr).to_string();
}
Ok(text)
}
fn xml_escape(raw: &str) -> String {
raw.replace('&', "&amp;")
.replace('<', "&lt;")
.replace('>', "&gt;")
.replace('"', "&quot;")
.replace('\'', "&apos;")
}
// ── Tests ───────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use crate::service::*;
#[test]
fn xml_escape_handles_reserved_chars() {
let escaped = xml_escape("<&>\"' and text");
assert_eq!(escaped, "&lt;&amp;&gt;&quot;&apos; and text");
}
#[test]
fn xml_escape_passes_through_plain_text() {
assert_eq!(xml_escape("hello world"), "hello world");
}
#[test]
fn run_capture_reads_stdout() {
let out = run_capture(Command::new("sh").args(["-c", "echo hello"]))
.expect("stdout capture should succeed");
assert_eq!(out.trim(), "hello");
}
#[test]
fn run_capture_falls_back_to_stderr() {
let out = run_capture(Command::new("sh").args(["-c", "echo warn 1>&2"]))
.expect("stderr capture should succeed");
assert_eq!(out.trim(), "warn");
}
#[test]
fn run_checked_errors_on_non_zero_exit() {
let err = run_checked(Command::new("sh").args(["-c", "exit 17"]))
.expect_err("non-zero exit should error");
assert!(err.to_string().contains("command failed"));
}
#[test]
fn run_checked_succeeds_on_zero_exit() {
assert!(run_checked(Command::new("sh").args(["-c", "exit 0"])).is_ok());
}
#[cfg(target_os = "macos")]
#[test]
fn macos_plist_path_has_expected_suffix() {
let path = macos_plist_path().unwrap();
let s = path.to_string_lossy();
assert!(
s.ends_with("Library/LaunchAgents/com.ironclaw.daemon.plist"),
"unexpected path: {s}"
);
}
#[cfg(target_os = "linux")]
#[test]
fn linux_unit_path_has_expected_suffix() {
let path = linux_unit_path().unwrap();
let s = path.to_string_lossy();
assert!(
s.ends_with(".config/systemd/user/ironclaw.service"),
"unexpected path: {s}"
);
}
#[test]
fn logs_dir_under_ironclaw() {
let path = ironclaw_logs_dir().unwrap();
let s = path.to_string_lossy();
assert!(s.ends_with(".ironclaw/logs"), "unexpected path: {s}");
}
}
+298 -1
View File
@@ -149,11 +149,52 @@ impl Default for EmbeddingsSettings {
/// Tunnel settings for public webhook endpoints.
///
/// The tunnel URL is shared across all channels that need webhooks.
/// Two modes:
/// - **Static URL**: `public_url` set directly (manual tunnel management).
/// - **Managed provider**: `provider` is set and the agent starts/stops the
/// tunnel process automatically at boot/shutdown.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TunnelSettings {
/// Public URL from tunnel provider (e.g., "https://abc123.ngrok.io").
/// When set without a provider, treated as a static (externally managed) URL.
#[serde(default)]
pub public_url: Option<String>,
/// Managed tunnel provider: "ngrok", "cloudflare", "tailscale", "custom".
#[serde(default)]
pub provider: Option<String>,
/// Cloudflare tunnel token.
#[serde(default)]
pub cf_token: Option<String>,
/// ngrok auth token.
#[serde(default)]
pub ngrok_token: Option<String>,
/// ngrok custom domain (paid plans).
#[serde(default)]
pub ngrok_domain: Option<String>,
/// Use Tailscale Funnel (public) instead of Serve (tailnet-only).
#[serde(default)]
pub ts_funnel: bool,
/// Tailscale hostname override.
#[serde(default)]
pub ts_hostname: Option<String>,
/// Shell command for custom tunnel (with `{port}` / `{host}` placeholders).
#[serde(default)]
pub custom_command: Option<String>,
/// Health check URL for custom tunnel.
#[serde(default)]
pub custom_health_url: Option<String>,
/// Substring pattern to extract URL from custom tunnel stdout.
#[serde(default)]
pub custom_url_pattern: Option<String>,
}
/// Channel-specific settings.
@@ -585,6 +626,83 @@ impl Settings {
}
}
/// Default TOML config file path (~/.ironclaw/config.toml).
pub fn default_toml_path() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw")
.join("config.toml")
}
/// Load settings from a TOML file.
///
/// Returns `None` if the file doesn't exist. Returns an error only
/// if the file exists but can't be parsed.
pub fn load_toml(path: &std::path::Path) -> Result<Option<Self>, String> {
let data = match std::fs::read_to_string(path) {
Ok(d) => d,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(e) => return Err(format!("failed to read {}: {}", path.display(), e)),
};
let settings: Self = toml::from_str(&data)
.map_err(|e| format!("invalid TOML in {}: {}", path.display(), e))?;
Ok(Some(settings))
}
/// Write a well-commented TOML config file with current settings.
pub fn save_toml(&self, path: &std::path::Path) -> Result<(), String> {
let raw = toml::to_string_pretty(self)
.map_err(|e| format!("failed to serialize settings: {}", e))?;
let content = format!(
"# IronClaw configuration file.\n\
#\n\
# Priority: env var > this file > database settings > defaults.\n\
# Uncomment and edit values to override defaults.\n\
# Run `ironclaw config init` to regenerate this file.\n\
#\n\
# Documentation: https://github.com/nearai/ironclaw\n\
\n\
{raw}"
);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| format!("failed to create {}: {}", parent.display(), e))?;
}
std::fs::write(path, content)
.map_err(|e| format!("failed to write {}: {}", path.display(), e))
}
/// Merge values from `other` into `self`, preferring `other` for
/// fields that differ from the default.
///
/// This enables layering: load DB/JSON settings as the base, then
/// overlay TOML values on top. Only fields that the TOML file
/// explicitly changed (i.e. differ from Default) are applied.
pub fn merge_from(&mut self, other: &Self) {
let default_json = match serde_json::to_value(Self::default()) {
Ok(v) => v,
Err(_) => return,
};
let other_json = match serde_json::to_value(other) {
Ok(v) => v,
Err(_) => return,
};
let mut self_json = match serde_json::to_value(&*self) {
Ok(v) => v,
Err(_) => return,
};
merge_non_default(&mut self_json, &other_json, &default_json);
if let Ok(merged) = serde_json::from_value(self_json) {
*self = merged;
}
}
/// Get a setting value by dotted path (e.g., "agent.max_parallel_jobs").
pub fn get(&self, path: &str) -> Option<String> {
let json = serde_json::to_value(self).ok()?;
@@ -766,9 +884,40 @@ fn collect_settings(
}
}
/// Recursively merge `other` into `target`, but only for fields where
/// `other` differs from `defaults`. This means only explicitly-set values
/// in the TOML file override the base settings.
fn merge_non_default(
target: &mut serde_json::Value,
other: &serde_json::Value,
defaults: &serde_json::Value,
) {
match (target, other, defaults) {
(
serde_json::Value::Object(t),
serde_json::Value::Object(o),
serde_json::Value::Object(d),
) => {
for (key, other_val) in o {
let default_val = d.get(key).cloned().unwrap_or(serde_json::Value::Null);
if let Some(target_val) = t.get_mut(key) {
merge_non_default(target_val, other_val, &default_val);
} else if other_val != &default_val {
t.insert(key.clone(), other_val.clone());
}
}
}
(target, other, defaults) => {
if other != defaults {
*target = other.clone();
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::settings::*;
#[test]
fn test_db_map_round_trip() {
@@ -904,4 +1053,152 @@ mod tests {
Some("http://my-vllm:8000/v1".to_string())
);
}
#[test]
fn toml_round_trip() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("config.toml");
let mut settings = Settings::default();
settings.agent.name = "toml-bot".to_string();
settings.heartbeat.enabled = true;
settings.heartbeat.interval_secs = 900;
settings.save_toml(&path).unwrap();
let loaded = Settings::load_toml(&path).unwrap().unwrap();
assert_eq!(loaded.agent.name, "toml-bot");
assert!(loaded.heartbeat.enabled);
assert_eq!(loaded.heartbeat.interval_secs, 900);
}
#[test]
fn toml_missing_file_returns_none() {
let result = Settings::load_toml(std::path::Path::new("/tmp/nonexistent_config.toml"));
assert!(result.unwrap().is_none());
}
#[test]
fn toml_invalid_content_returns_error() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("bad.toml");
std::fs::write(&path, "this is not valid toml [[[").unwrap();
let result = Settings::load_toml(&path);
assert!(result.is_err());
}
#[test]
fn toml_partial_config_uses_defaults() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("partial.toml");
// Only set agent name, everything else should be default
std::fs::write(&path, "[agent]\nname = \"partial-bot\"\n").unwrap();
let loaded = Settings::load_toml(&path).unwrap().unwrap();
assert_eq!(loaded.agent.name, "partial-bot");
// Defaults preserved
assert_eq!(loaded.agent.max_parallel_jobs, 5);
assert!(!loaded.heartbeat.enabled);
}
#[test]
fn toml_header_comment_present() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("config.toml");
Settings::default().save_toml(&path).unwrap();
let content = std::fs::read_to_string(&path).unwrap();
assert!(content.starts_with("# IronClaw configuration file."));
assert!(content.contains("[agent]"));
assert!(content.contains("[heartbeat]"));
}
#[test]
fn merge_only_overrides_non_default_values() {
let mut base = Settings::default();
base.agent.name = "from-db".to_string();
base.heartbeat.interval_secs = 600;
let mut toml_overlay = Settings::default();
toml_overlay.agent.name = "from-toml".to_string();
// heartbeat.interval_secs stays at default (1800) in the overlay,
// so the base value (600) should be preserved.
base.merge_from(&toml_overlay);
assert_eq!(base.agent.name, "from-toml");
assert_eq!(base.heartbeat.interval_secs, 600);
}
#[test]
fn merge_preserves_base_when_overlay_is_default() {
let mut base = Settings::default();
base.agent.name = "custom-name".to_string();
base.heartbeat.enabled = true;
let overlay = Settings::default();
base.merge_from(&overlay);
// All base values preserved since overlay is entirely default
assert_eq!(base.agent.name, "custom-name");
assert!(base.heartbeat.enabled);
}
#[test]
fn toml_creates_parent_dirs() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("nested").join("deep").join("config.toml");
Settings::default().save_toml(&path).unwrap();
assert!(path.exists());
}
#[test]
fn default_toml_path_under_ironclaw() {
let path = Settings::default_toml_path();
assert!(path.to_string_lossy().contains(".ironclaw"));
assert!(path.to_string_lossy().ends_with("config.toml"));
}
#[test]
fn tunnel_settings_round_trip() {
let settings = Settings {
tunnel: TunnelSettings {
provider: Some("ngrok".to_string()),
ngrok_token: Some("tok_abc123".to_string()),
ngrok_domain: Some("my.ngrok.dev".to_string()),
..Default::default()
},
..Default::default()
};
// JSON round-trip
let json = serde_json::to_string(&settings).unwrap();
let restored: Settings = serde_json::from_str(&json).unwrap();
assert_eq!(restored.tunnel.provider, Some("ngrok".to_string()));
assert_eq!(restored.tunnel.ngrok_token, Some("tok_abc123".to_string()));
assert_eq!(
restored.tunnel.ngrok_domain,
Some("my.ngrok.dev".to_string())
);
assert!(restored.tunnel.public_url.is_none());
// DB map round-trip
let map = settings.to_db_map();
let from_db = Settings::from_db_map(&map);
assert_eq!(from_db.tunnel.provider, Some("ngrok".to_string()));
assert_eq!(from_db.tunnel.ngrok_token, Some("tok_abc123".to_string()));
// get/set round-trip
let mut s = Settings::default();
s.set("tunnel.provider", "cloudflare").unwrap();
s.set("tunnel.cf_token", "cf_tok_xyz").unwrap();
s.set("tunnel.ts_funnel", "true").unwrap();
assert_eq!(s.tunnel.provider, Some("cloudflare".to_string()));
assert_eq!(s.tunnel.cf_token, Some("cf_tok_xyz".to_string()));
assert!(s.tunnel.ts_funnel);
}
}
+127 -21
View File
@@ -18,6 +18,7 @@ use crate::secrets::{CreateSecretParams, SecretsStore};
use crate::settings::{Settings, TunnelSettings};
use crate::setup::prompts::{
confirm, input, optional_input, print_error, print_info, print_success, secret_input,
select_one,
};
/// Typed errors for channel setup flows.
@@ -356,12 +357,20 @@ async fn bind_telegram_owner_flow(
/// Set up a tunnel for exposing the agent to the internet.
///
/// This is shared across all channels that need webhook endpoints.
/// Returns the tunnel URL if configured.
pub fn setup_tunnel(settings: &Settings) -> Result<Option<String>, ChannelSetupError> {
if let Some(ref url) = settings.tunnel.public_url {
print_info(&format!("Existing tunnel configured: {}", url));
/// Returns a `TunnelSettings` with provider config (managed tunnel)
/// or a static URL.
pub fn setup_tunnel(settings: &Settings) -> Result<TunnelSettings, ChannelSetupError> {
// Show existing config
let has_existing = settings.tunnel.public_url.is_some() || settings.tunnel.provider.is_some();
if has_existing {
if let Some(ref url) = settings.tunnel.public_url {
print_info(&format!("Existing static tunnel URL: {}", url));
}
if let Some(ref provider) = settings.tunnel.provider {
print_info(&format!("Existing managed provider: {}", provider));
}
if !confirm("Change tunnel configuration?", false)? {
return Ok(Some(url.clone()));
return Ok(settings.tunnel.clone());
}
}
@@ -369,24 +378,121 @@ pub fn setup_tunnel(settings: &Settings) -> Result<Option<String>, ChannelSetupE
print_info("Tunnel Configuration (for webhook endpoints):");
print_info("A tunnel exposes your local agent to the internet, enabling:");
print_info(" - Instant Telegram message delivery (instead of polling)");
print_info(" - Future: Slack, Discord, GitHub webhooks");
print_info("");
print_info("Supported tunnel providers:");
print_info(" - ngrok: ngrok http 8080");
print_info(" - Cloudflare: cloudflared tunnel --url http://localhost:8080");
print_info(" - localtunnel: lt --port 8080");
print_info("");
print_info("Security note: Webhook endpoints don't use tunnel-level auth.");
print_info("Security comes from provider-specific secrets (e.g., Telegram webhook secret).");
print_info(" - Slack, Discord, GitHub webhooks");
println!();
if !confirm("Configure a tunnel?", false)? {
return Ok(None);
return Ok(TunnelSettings::default());
}
let options = &[
"ngrok - managed tunnel, starts automatically",
"Cloudflare - cloudflared tunnel, starts automatically",
"Tailscale - Tailscale Funnel/Serve, starts automatically",
"Custom - your own tunnel command",
"Static URL - you manage the tunnel yourself",
];
let choice = select_one("Select tunnel provider:", options)?;
match choice {
0 => setup_tunnel_ngrok(),
1 => setup_tunnel_cloudflare(),
2 => setup_tunnel_tailscale(),
3 => setup_tunnel_custom(),
4 => setup_tunnel_static(),
_ => Ok(TunnelSettings::default()),
}
}
fn setup_tunnel_ngrok() -> Result<TunnelSettings, ChannelSetupError> {
print_info("Get your auth token from: https://dashboard.ngrok.com/get-started/your-authtoken");
println!();
let token = secret_input("ngrok auth token")?;
let domain = optional_input("Custom domain", Some("leave empty for auto-assigned"))?;
print_success("ngrok configured. Tunnel will start automatically at boot.");
Ok(TunnelSettings {
provider: Some("ngrok".to_string()),
ngrok_token: Some(token.expose_secret().to_string()),
ngrok_domain: domain,
..Default::default()
})
}
fn setup_tunnel_cloudflare() -> Result<TunnelSettings, ChannelSetupError> {
print_info("Get your tunnel token from the Cloudflare Zero Trust dashboard:");
print_info(" https://one.dash.cloudflare.com/ > Networks > Tunnels");
println!();
let token = secret_input("Cloudflare tunnel token")?;
print_success("Cloudflare tunnel configured. Tunnel will start automatically at boot.");
Ok(TunnelSettings {
provider: Some("cloudflare".to_string()),
cf_token: Some(token.expose_secret().to_string()),
..Default::default()
})
}
fn setup_tunnel_tailscale() -> Result<TunnelSettings, ChannelSetupError> {
let funnel = confirm("Use Tailscale Funnel (public internet)?", true)?;
let hostname = optional_input("Hostname override", Some("leave empty for auto-detect"))?;
let mode = if funnel {
"Funnel (public)"
} else {
"Serve (tailnet-only)"
};
print_success(&format!("Tailscale {} configured.", mode));
Ok(TunnelSettings {
provider: Some("tailscale".to_string()),
ts_funnel: funnel,
ts_hostname: hostname,
..Default::default()
})
}
fn setup_tunnel_custom() -> Result<TunnelSettings, ChannelSetupError> {
print_info("Enter a shell command to start your tunnel.");
print_info("Use {port} and {host} as placeholders.");
print_info("Example: bore local {port} --to bore.pub");
println!();
let command = input("Tunnel command")?;
if command.is_empty() {
return Err(ChannelSetupError::Validation(
"Tunnel command cannot be empty".to_string(),
));
}
let health_url = optional_input("Health check URL", Some("optional"))?;
let url_pattern = optional_input(
"URL pattern (substring to match in stdout)",
Some("optional"),
)?;
print_success("Custom tunnel configured.");
Ok(TunnelSettings {
provider: Some("custom".to_string()),
custom_command: Some(command),
custom_health_url: health_url,
custom_url_pattern: url_pattern,
..Default::default()
})
}
fn setup_tunnel_static() -> Result<TunnelSettings, ChannelSetupError> {
print_info("Enter the public URL of your externally managed tunnel.");
println!();
let tunnel_url = input("Tunnel URL (e.g., https://abc123.ngrok.io)")?;
// Validate URL format
if !tunnel_url.starts_with("https://") {
print_error("URL must start with https:// (webhooks require HTTPS)");
return Err(ChannelSetupError::Validation(
@@ -394,15 +500,15 @@ pub fn setup_tunnel(settings: &Settings) -> Result<Option<String>, ChannelSetupE
));
}
// Remove trailing slash if present
let tunnel_url = tunnel_url.trim_end_matches('/').to_string();
print_success(&format!("Tunnel URL configured: {}", tunnel_url));
print_info("");
print_success(&format!("Static tunnel URL configured: {}", tunnel_url));
print_info("Make sure your tunnel is running before starting the agent.");
print_info("You can also set TUNNEL_URL environment variable to override.");
Ok(Some(tunnel_url))
Ok(TunnelSettings {
public_url: Some(tunnel_url),
..Default::default()
})
}
/// Set up Telegram webhook secret for signature validation.
+14 -11
View File
@@ -1022,6 +1022,11 @@ impl SetupWizard {
api_key: None,
fallback_model: None,
max_retries: 3,
circuit_breaker_threshold: None,
circuit_breaker_recovery_secs: 30,
response_cache_enabled: false,
response_cache_ttl_secs: 3600,
response_cache_max_entries: 1000,
failover_cooldown_secs: 300,
failover_cooldown_threshold: 3,
},
@@ -1254,11 +1259,8 @@ impl SetupWizard {
async fn step_channels(&mut self) -> Result<(), SetupError> {
// First, configure tunnel (shared across all channels that need webhooks)
match setup_tunnel(&self.settings) {
Ok(Some(url)) => {
self.settings.tunnel.public_url = Some(url);
}
Ok(None) => {
self.settings.tunnel.public_url = None;
Ok(tunnel_settings) => {
self.settings.tunnel = tunnel_settings;
}
Err(e) => {
print_info(&format!("Tunnel setup skipped: {}", e));
@@ -1609,9 +1611,14 @@ impl SetupWizard {
}
if let Some(ref tunnel_url) = self.settings.tunnel.public_url {
println!(" Tunnel: {}", tunnel_url);
println!(" Tunnel: {} (static)", tunnel_url);
} else if let Some(ref provider) = self.settings.tunnel.provider {
println!(" Tunnel: {} (managed, starts at boot)", provider);
}
let has_tunnel =
self.settings.tunnel.public_url.is_some() || self.settings.tunnel.provider.is_some();
println!(" Channels:");
println!(" - CLI/TUI: enabled");
@@ -1621,11 +1628,7 @@ impl SetupWizard {
}
for channel_name in &self.settings.channels.wasm_channels {
let mode = if self.settings.tunnel.public_url.is_some() {
"webhook"
} else {
"polling"
};
let mode = if has_tunnel { "webhook" } else { "polling" };
println!(
" - {}: enabled ({})",
capitalize_first(channel_name),
+140
View File
@@ -0,0 +1,140 @@
//! Cloudflare Tunnel via the `cloudflared` binary.
use anyhow::{Result, bail};
use tokio::io::AsyncBufReadExt;
use tokio::process::Command;
use crate::tunnel::{
SharedProcess, SharedUrl, Tunnel, TunnelProcess, kill_shared, new_shared_process,
new_shared_url,
};
/// Wraps `cloudflared` with token-based auth from the Zero Trust dashboard.
pub struct CloudflareTunnel {
token: String,
proc: SharedProcess,
url: SharedUrl,
}
impl CloudflareTunnel {
pub fn new(token: String) -> Self {
Self {
token,
proc: new_shared_process(),
url: new_shared_url(),
}
}
}
#[async_trait::async_trait]
impl Tunnel for CloudflareTunnel {
fn name(&self) -> &str {
"cloudflare"
}
async fn start(&self, local_host: &str, local_port: u16) -> Result<String> {
let origin = format!("http://{local_host}:{local_port}");
let mut child = Command::new("cloudflared")
.args([
"tunnel",
"--no-autoupdate",
"run",
"--token",
&self.token,
"--url",
&origin,
])
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true)
.spawn()?;
// cloudflared prints the public URL on stderr
let stderr = child
.stderr
.take()
.ok_or_else(|| anyhow::anyhow!("Failed to capture cloudflared stderr"))?;
let mut reader = tokio::io::BufReader::new(stderr).lines();
let mut public_url = String::new();
let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(30);
while tokio::time::Instant::now() < deadline {
let line =
tokio::time::timeout(tokio::time::Duration::from_secs(5), reader.next_line()).await;
match line {
Ok(Ok(Some(l))) => {
tracing::debug!("cloudflared: {l}");
if let Some(idx) = l.find("https://") {
let url_part = &l[idx..];
let end = url_part
.find(|c: char| c.is_whitespace())
.unwrap_or(url_part.len());
public_url = url_part[..end].to_string();
break;
}
}
Ok(Ok(None)) => break,
Ok(Err(e)) => bail!("Error reading cloudflared output: {e}"),
Err(_) => {} // line timeout, keep waiting
}
}
if public_url.is_empty() {
child.kill().await.ok();
bail!("cloudflared did not produce a public URL within 30s. Is the token valid?");
}
if let Ok(mut guard) = self.url.write() {
*guard = Some(public_url.clone());
}
let mut guard = self.proc.lock().await;
*guard = Some(TunnelProcess { child });
Ok(public_url)
}
async fn stop(&self) -> Result<()> {
if let Ok(mut guard) = self.url.write() {
*guard = None;
}
kill_shared(&self.proc).await
}
async fn health_check(&self) -> bool {
let guard = self.proc.lock().await;
guard.as_ref().is_some_and(|tp| tp.child.id().is_some())
}
fn public_url(&self) -> Option<String> {
self.url.read().ok().and_then(|guard| guard.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn constructor_stores_token() {
let tunnel = CloudflareTunnel::new("cf-token".into());
assert_eq!(tunnel.token, "cf-token");
}
#[test]
fn public_url_none_before_start() {
assert!(CloudflareTunnel::new("tok".into()).public_url().is_none());
}
#[tokio::test]
async fn stop_without_start_is_ok() {
assert!(CloudflareTunnel::new("tok".into()).stop().await.is_ok());
}
#[tokio::test]
async fn health_false_before_start() {
assert!(!CloudflareTunnel::new("tok".into()).health_check().await);
}
}
+245
View File
@@ -0,0 +1,245 @@
//! Custom tunnel via an arbitrary shell command.
use anyhow::{Result, bail};
use tokio::io::AsyncBufReadExt;
use tokio::process::Command;
use crate::tunnel::{
SharedProcess, SharedUrl, Tunnel, TunnelProcess, kill_shared, new_shared_process,
new_shared_url,
};
/// Bring-your-own tunnel binary.
///
/// `start_command` supports `{port}` and `{host}` placeholders.
/// If `url_pattern` is set, stdout is scanned for a URL matching that
/// substring. If `health_url` is set, health checks poll that endpoint.
///
/// **Note:** The command is split on whitespace, so quoted arguments like
/// `--arg "hello world"` won't work. Each token must be a single word.
///
/// Examples:
/// - `bore local {port} --to bore.pub`
/// - `ssh -R 80:localhost:{port} serveo.net`
pub struct CustomTunnel {
start_command: String,
health_url: Option<String>,
url_pattern: Option<String>,
proc: SharedProcess,
url: SharedUrl,
}
impl CustomTunnel {
pub fn new(
start_command: String,
health_url: Option<String>,
url_pattern: Option<String>,
) -> Self {
Self {
start_command,
health_url,
url_pattern,
proc: new_shared_process(),
url: new_shared_url(),
}
}
}
#[async_trait::async_trait]
impl Tunnel for CustomTunnel {
fn name(&self) -> &str {
"custom"
}
async fn start(&self, local_host: &str, local_port: u16) -> Result<String> {
let cmd = self
.start_command
.replace("{port}", &local_port.to_string())
.replace("{host}", local_host);
let parts: Vec<&str> = cmd.split_whitespace().collect();
if parts.is_empty() {
bail!("Custom tunnel start_command is empty");
}
let mut child = Command::new(parts[0])
.args(&parts[1..])
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true)
.spawn()?;
let mut public_url = format!("http://{local_host}:{local_port}");
if self.url_pattern.is_some()
&& let Some(stdout) = child.stdout.take()
{
let mut reader = tokio::io::BufReader::new(stdout).lines();
let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(15);
while tokio::time::Instant::now() < deadline {
let line =
tokio::time::timeout(tokio::time::Duration::from_secs(3), reader.next_line())
.await;
match line {
Ok(Ok(Some(l))) => {
tracing::debug!("custom-tunnel: {l}");
if let Some(url) = extract_url(&l) {
let matches_pattern = self
.url_pattern
.as_ref()
.is_none_or(|pat| url.contains(pat.as_str()));
if matches_pattern {
public_url = url;
break;
}
}
}
Ok(Ok(None) | Err(_)) => break,
Err(_) => {}
}
}
}
if let Ok(mut guard) = self.url.write() {
*guard = Some(public_url.clone());
}
let mut guard = self.proc.lock().await;
*guard = Some(TunnelProcess { child });
Ok(public_url)
}
async fn stop(&self) -> Result<()> {
if let Ok(mut guard) = self.url.write() {
*guard = None;
}
kill_shared(&self.proc).await
}
async fn health_check(&self) -> bool {
if let Some(ref url) = self.health_url {
return reqwest::Client::new()
.get(url)
.timeout(std::time::Duration::from_secs(5))
.send()
.await
.is_ok();
}
let guard = self.proc.lock().await;
guard.as_ref().is_some_and(|tp| tp.child.id().is_some())
}
fn public_url(&self) -> Option<String> {
self.url.read().ok().and_then(|guard| guard.clone())
}
}
/// Extract the first `https://` or `http://` URL from a line of text.
fn extract_url(line: &str) -> Option<String> {
let idx = line.find("https://").or_else(|| line.find("http://"))?;
let url_part = &line[idx..];
let end = url_part
.find(|c: char| c.is_whitespace())
.unwrap_or(url_part.len());
Some(url_part[..end].to_string())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn empty_command_returns_error() {
let tunnel = CustomTunnel::new(" ".into(), None, None);
let result = tunnel.start("127.0.0.1", 8080).await;
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("start_command is empty")
);
}
#[tokio::test]
async fn start_without_pattern_returns_local() {
let tunnel = CustomTunnel::new("sleep 1".into(), None, None);
let url = tunnel.start("127.0.0.1", 4455).await.unwrap();
assert_eq!(url, "http://127.0.0.1:4455");
tunnel.stop().await.unwrap();
}
#[tokio::test]
async fn start_with_pattern_extracts_url() {
let tunnel = CustomTunnel::new(
"echo https://public.example".into(),
None,
Some("public.example".into()),
);
let url = tunnel.start("localhost", 9999).await.unwrap();
assert_eq!(url, "https://public.example");
tunnel.stop().await.unwrap();
}
#[tokio::test]
async fn pattern_filters_non_matching_urls() {
// The command outputs two lines: first a non-matching URL, then a matching one.
// The pattern filter should skip the first and grab the second.
// No shell quoting needed; Command passes args directly to the binary.
let tunnel = CustomTunnel::new(
r"printf http://internal:1234\nhttps://real.tunnel.io/abc\n".into(),
None,
Some("tunnel.io".into()),
);
let url = tunnel.start("localhost", 9999).await.unwrap();
assert_eq!(url, "https://real.tunnel.io/abc");
tunnel.stop().await.unwrap();
}
#[tokio::test]
async fn replaces_host_and_port_placeholders() {
let tunnel = CustomTunnel::new(
"echo http://{host}:{port}".into(),
None,
Some("http://".into()),
);
let url = tunnel.start("10.1.2.3", 4321).await.unwrap();
assert_eq!(url, "http://10.1.2.3:4321");
tunnel.stop().await.unwrap();
}
#[tokio::test]
async fn health_with_unreachable_url_is_false() {
let tunnel = CustomTunnel::new(
"sleep 1".into(),
Some("http://127.0.0.1:9/healthz".into()),
None,
);
assert!(!tunnel.health_check().await);
}
#[test]
fn extract_url_finds_https() {
assert_eq!(
extract_url("tunnel ready at https://foo.bar.com/path more text"),
Some("https://foo.bar.com/path".to_string())
);
}
#[test]
fn extract_url_finds_http() {
assert_eq!(
extract_url("url=http://localhost:8080"),
Some("http://localhost:8080".to_string())
);
}
#[test]
fn extract_url_none_when_absent() {
assert_eq!(extract_url("no url here"), None);
}
}
+329
View File
@@ -0,0 +1,329 @@
//! Tunnel abstraction for exposing the agent to the internet.
//!
//! Wraps external tunnel binaries (cloudflared, ngrok, tailscale, etc.) behind
//! a common trait. The gateway starts a tunnel after binding its local port
//! and stops it on shutdown.
//!
//! Supported providers:
//! - **cloudflare** - Zero Trust tunnels via `cloudflared`
//! - **tailscale** - `tailscale serve` (tailnet) or `tailscale funnel` (public)
//! - **ngrok** - instant public URLs via `ngrok`
//! - **custom** - any command with `{host}`/`{port}` placeholders
//! - **none** - local-only, no external exposure
mod cloudflare;
mod custom;
mod ngrok;
mod none;
mod tailscale;
pub use cloudflare::CloudflareTunnel;
pub use custom::CustomTunnel;
pub use ngrok::NgrokTunnel;
pub use none::NoneTunnel;
pub use tailscale::TailscaleTunnel;
use std::sync::Arc;
use anyhow::{Result, bail};
use tokio::sync::Mutex;
/// Lock-free URL storage. Uses `std::sync::RwLock` so `public_url()` (sync)
/// never returns a spurious `None` due to async lock contention.
pub(crate) type SharedUrl = Arc<std::sync::RwLock<Option<String>>>;
pub(crate) fn new_shared_url() -> SharedUrl {
Arc::new(std::sync::RwLock::new(None))
}
// ── Tunnel trait ─────────────────────────────────────────────────
/// Provider-agnostic tunnel with lifecycle management.
///
/// Implementations wrap an external tunnel binary. The gateway calls
/// `start()` after binding its local port and `stop()` on shutdown.
#[async_trait::async_trait]
pub trait Tunnel: Send + Sync {
/// Human-readable provider name (e.g. "cloudflare", "tailscale").
fn name(&self) -> &str;
/// Start the tunnel exposing `local_host:local_port` externally.
/// Returns the public URL on success.
async fn start(&self, local_host: &str, local_port: u16) -> Result<String>;
/// Stop the tunnel process gracefully.
async fn stop(&self) -> Result<()>;
/// Check if the tunnel process is still alive.
async fn health_check(&self) -> bool;
/// Return the public URL if the tunnel is running, `None` otherwise.
fn public_url(&self) -> Option<String>;
}
// ── Shared child-process handle ──────────────────────────────────
/// Wraps a spawned tunnel child process.
pub(crate) struct TunnelProcess {
pub child: tokio::process::Child,
}
pub(crate) type SharedProcess = Arc<Mutex<Option<TunnelProcess>>>;
pub(crate) fn new_shared_process() -> SharedProcess {
Arc::new(Mutex::new(None))
}
/// Kill a shared tunnel process if running.
pub(crate) async fn kill_shared(proc: &SharedProcess) -> Result<()> {
let mut guard = proc.lock().await;
if let Some(ref mut tp) = *guard {
tp.child.kill().await.ok();
tp.child.wait().await.ok();
}
*guard = None;
Ok(())
}
// ── Configuration types ──────────────────────────────────────────
/// Provider-specific config for Cloudflare tunnels.
#[derive(Debug, Clone, Default)]
pub struct CloudflareTunnelConfig {
/// Token from the Cloudflare Zero Trust dashboard.
pub token: String,
}
/// Provider-specific config for Tailscale tunnels.
#[derive(Debug, Clone, Default)]
pub struct TailscaleTunnelConfig {
/// Use `tailscale funnel` (public) instead of `tailscale serve` (tailnet).
pub funnel: bool,
/// Override the hostname (default: auto-detect from `tailscale status`).
pub hostname: Option<String>,
}
/// Provider-specific config for ngrok tunnels.
#[derive(Debug, Clone, Default)]
pub struct NgrokTunnelConfig {
/// ngrok auth token (required).
pub auth_token: String,
/// Custom domain (requires ngrok paid plan).
pub domain: Option<String>,
}
/// Provider-specific config for custom tunnel commands.
#[derive(Debug, Clone, Default)]
pub struct CustomTunnelConfig {
/// Shell command with `{port}` and `{host}` placeholders.
pub start_command: String,
/// HTTP endpoint to poll for health checks.
pub health_url: Option<String>,
/// Substring to match in stdout for URL extraction.
pub url_pattern: Option<String>,
}
/// Full tunnel configuration.
#[derive(Debug, Clone, Default)]
pub struct TunnelProviderConfig {
/// Provider name: "none", "cloudflare", "tailscale", "ngrok", "custom".
pub provider: String,
pub cloudflare: Option<CloudflareTunnelConfig>,
pub tailscale: Option<TailscaleTunnelConfig>,
pub ngrok: Option<NgrokTunnelConfig>,
pub custom: Option<CustomTunnelConfig>,
}
// ── Factory ──────────────────────────────────────────────────────
/// Create a tunnel from config. Returns `None` for provider "none" or empty.
pub fn create_tunnel(config: &TunnelProviderConfig) -> Result<Option<Box<dyn Tunnel>>> {
match config.provider.as_str() {
"none" | "" => Ok(None),
"cloudflare" => {
let cf = config.cloudflare.as_ref().ok_or_else(|| {
anyhow::anyhow!("TUNNEL_PROVIDER=cloudflare but no TUNNEL_CF_TOKEN configured")
})?;
Ok(Some(Box::new(CloudflareTunnel::new(cf.token.clone()))))
}
"tailscale" => {
let ts = config.tailscale.as_ref().cloned().unwrap_or_default();
Ok(Some(Box::new(TailscaleTunnel::new(ts.funnel, ts.hostname))))
}
"ngrok" => {
let ng = config.ngrok.as_ref().ok_or_else(|| {
anyhow::anyhow!("TUNNEL_PROVIDER=ngrok but no TUNNEL_NGROK_TOKEN configured")
})?;
Ok(Some(Box::new(NgrokTunnel::new(
ng.auth_token.clone(),
ng.domain.clone(),
))))
}
"custom" => {
let cu = config.custom.as_ref().ok_or_else(|| {
anyhow::anyhow!("TUNNEL_PROVIDER=custom but no TUNNEL_CUSTOM_COMMAND configured")
})?;
Ok(Some(Box::new(CustomTunnel::new(
cu.start_command.clone(),
cu.health_url.clone(),
cu.url_pattern.clone(),
))))
}
other => bail!(
"Unknown tunnel provider: \"{other}\". Valid: none, cloudflare, tailscale, ngrok, custom"
),
}
}
// ── Tests ────────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
use super::*;
use tokio::process::Command;
fn assert_tunnel_err(cfg: &TunnelProviderConfig, needle: &str) {
match create_tunnel(cfg) {
Err(e) => assert!(
e.to_string().contains(needle),
"Expected error containing \"{needle}\", got: {e}"
),
Ok(_) => panic!("Expected error containing \"{needle}\", but got Ok"),
}
}
#[test]
fn factory_none_returns_none() {
let cfg = TunnelProviderConfig::default();
assert!(create_tunnel(&cfg).unwrap().is_none());
}
#[test]
fn factory_empty_returns_none() {
let cfg = TunnelProviderConfig {
provider: String::new(),
..Default::default()
};
assert!(create_tunnel(&cfg).unwrap().is_none());
}
#[test]
fn factory_unknown_provider_errors() {
let cfg = TunnelProviderConfig {
provider: "wireguard".into(),
..Default::default()
};
assert_tunnel_err(&cfg, "Unknown tunnel provider");
}
#[test]
fn factory_cloudflare_missing_config_errors() {
let cfg = TunnelProviderConfig {
provider: "cloudflare".into(),
..Default::default()
};
assert_tunnel_err(&cfg, "TUNNEL_CF_TOKEN");
}
#[test]
fn factory_cloudflare_with_config_ok() {
let cfg = TunnelProviderConfig {
provider: "cloudflare".into(),
cloudflare: Some(CloudflareTunnelConfig {
token: "test-token".into(),
}),
..Default::default()
};
let t = create_tunnel(&cfg).unwrap().unwrap();
assert_eq!(t.name(), "cloudflare");
}
#[test]
fn factory_tailscale_defaults_ok() {
let cfg = TunnelProviderConfig {
provider: "tailscale".into(),
..Default::default()
};
let t = create_tunnel(&cfg).unwrap().unwrap();
assert_eq!(t.name(), "tailscale");
}
#[test]
fn factory_ngrok_missing_config_errors() {
let cfg = TunnelProviderConfig {
provider: "ngrok".into(),
..Default::default()
};
assert_tunnel_err(&cfg, "TUNNEL_NGROK_TOKEN");
}
#[test]
fn factory_ngrok_with_config_ok() {
let cfg = TunnelProviderConfig {
provider: "ngrok".into(),
ngrok: Some(NgrokTunnelConfig {
auth_token: "tok".into(),
domain: None,
}),
..Default::default()
};
let t = create_tunnel(&cfg).unwrap().unwrap();
assert_eq!(t.name(), "ngrok");
}
#[test]
fn factory_custom_missing_config_errors() {
let cfg = TunnelProviderConfig {
provider: "custom".into(),
..Default::default()
};
assert_tunnel_err(&cfg, "TUNNEL_CUSTOM_COMMAND");
}
#[test]
fn factory_custom_with_config_ok() {
let cfg = TunnelProviderConfig {
provider: "custom".into(),
custom: Some(CustomTunnelConfig {
start_command: "echo tunnel".into(),
health_url: None,
url_pattern: None,
}),
..Default::default()
};
let t = create_tunnel(&cfg).unwrap().unwrap();
assert_eq!(t.name(), "custom");
}
#[tokio::test]
async fn kill_shared_no_process_is_ok() {
let proc = new_shared_process();
assert!(kill_shared(&proc).await.is_ok());
assert!(proc.lock().await.is_none());
}
#[tokio::test]
async fn kill_shared_terminates_child() {
let proc = new_shared_process();
let child = Command::new("sleep")
.arg("30")
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.spawn()
.expect("sleep should spawn");
{
let mut guard = proc.lock().await;
*guard = Some(TunnelProcess { child });
}
kill_shared(&proc).await.unwrap();
assert!(proc.lock().await.is_none());
}
}
+142
View File
@@ -0,0 +1,142 @@
//! ngrok tunnel via the `ngrok` binary.
use anyhow::{Result, bail};
use tokio::io::AsyncBufReadExt;
use tokio::process::Command;
use crate::tunnel::{
SharedProcess, SharedUrl, Tunnel, TunnelProcess, kill_shared, new_shared_process,
new_shared_url,
};
/// Wraps `ngrok` with optional custom domain support (paid plan).
pub struct NgrokTunnel {
auth_token: String,
domain: Option<String>,
proc: SharedProcess,
url: SharedUrl,
}
impl NgrokTunnel {
pub fn new(auth_token: String, domain: Option<String>) -> Self {
Self {
auth_token,
domain,
proc: new_shared_process(),
url: new_shared_url(),
}
}
}
#[async_trait::async_trait]
impl Tunnel for NgrokTunnel {
fn name(&self) -> &str {
"ngrok"
}
async fn start(&self, local_host: &str, local_port: u16) -> Result<String> {
let mut args = vec!["http".to_string(), format!("{local_host}:{local_port}")];
if let Some(ref domain) = self.domain {
args.push("--domain".into());
args.push(domain.clone());
}
args.extend(["--log", "stdout", "--log-format", "logfmt"].map(String::from));
let mut child = Command::new("ngrok")
.args(&args)
.env("NGROK_AUTHTOKEN", &self.auth_token)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true)
.spawn()?;
let stdout = child
.stdout
.take()
.ok_or_else(|| anyhow::anyhow!("Failed to capture ngrok stdout"))?;
let mut reader = tokio::io::BufReader::new(stdout).lines();
let mut public_url = String::new();
let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(15);
while tokio::time::Instant::now() < deadline {
let line =
tokio::time::timeout(tokio::time::Duration::from_secs(3), reader.next_line()).await;
match line {
Ok(Ok(Some(l))) => {
tracing::debug!("ngrok: {l}");
// ngrok logfmt: url=https://xxxx.ngrok-free.app
if let Some(idx) = l.find("url=https://") {
let url_start = idx + 4; // skip "url="
let url_part = &l[url_start..];
let end = url_part
.find(|c: char| c.is_whitespace())
.unwrap_or(url_part.len());
public_url = url_part[..end].to_string();
break;
}
}
Ok(Ok(None)) => break,
Ok(Err(e)) => bail!("Error reading ngrok output: {e}"),
Err(_) => {}
}
}
if public_url.is_empty() {
child.kill().await.ok();
bail!("ngrok did not produce a public URL within 15s. Is the auth token valid?");
}
if let Ok(mut guard) = self.url.write() {
*guard = Some(public_url.clone());
}
let mut guard = self.proc.lock().await;
*guard = Some(TunnelProcess { child });
Ok(public_url)
}
async fn stop(&self) -> Result<()> {
if let Ok(mut guard) = self.url.write() {
*guard = None;
}
kill_shared(&self.proc).await
}
async fn health_check(&self) -> bool {
let guard = self.proc.lock().await;
guard.as_ref().is_some_and(|tp| tp.child.id().is_some())
}
fn public_url(&self) -> Option<String> {
self.url.read().ok().and_then(|guard| guard.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn constructor_stores_domain() {
let tunnel = NgrokTunnel::new("tok".into(), Some("my.ngrok.app".into()));
assert_eq!(tunnel.domain.as_deref(), Some("my.ngrok.app"));
}
#[test]
fn public_url_none_before_start() {
assert!(NgrokTunnel::new("tok".into(), None).public_url().is_none());
}
#[tokio::test]
async fn stop_without_start_is_ok() {
assert!(NgrokTunnel::new("tok".into(), None).stop().await.is_ok());
}
#[tokio::test]
async fn health_false_before_start() {
assert!(!NgrokTunnel::new("tok".into(), None).health_check().await);
}
}
+62
View File
@@ -0,0 +1,62 @@
//! No-op tunnel for local-only access.
use anyhow::Result;
use crate::tunnel::Tunnel;
/// No-op tunnel, no external exposure. `public_url()` always returns `None`.
pub struct NoneTunnel;
#[async_trait::async_trait]
impl Tunnel for NoneTunnel {
fn name(&self) -> &str {
"none"
}
async fn start(&self, local_host: &str, local_port: u16) -> Result<String> {
Ok(format!("http://{local_host}:{local_port}"))
}
async fn stop(&self) -> Result<()> {
Ok(())
}
async fn health_check(&self) -> bool {
true
}
fn public_url(&self) -> Option<String> {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn name_is_none() {
assert_eq!(NoneTunnel.name(), "none");
}
#[tokio::test]
async fn start_returns_local_url() {
let url = NoneTunnel.start("127.0.0.1", 7788).await.unwrap();
assert_eq!(url, "http://127.0.0.1:7788");
}
#[tokio::test]
async fn stop_is_noop() {
assert!(NoneTunnel.stop().await.is_ok());
}
#[tokio::test]
async fn health_is_always_true() {
assert!(NoneTunnel.health_check().await);
}
#[test]
fn public_url_is_always_none() {
assert!(NoneTunnel.public_url().is_none());
}
}
+140
View File
@@ -0,0 +1,140 @@
//! Tailscale tunnel via `tailscale serve` or `tailscale funnel`.
use anyhow::{Result, bail};
use tokio::process::Command;
use crate::tunnel::{
SharedProcess, SharedUrl, Tunnel, TunnelProcess, kill_shared, new_shared_process,
new_shared_url,
};
/// Uses `tailscale serve` (tailnet-only) or `tailscale funnel` (public).
///
/// Requires Tailscale installed and authenticated (`tailscale up`).
pub struct TailscaleTunnel {
funnel: bool,
hostname: Option<String>,
proc: SharedProcess,
url: SharedUrl,
}
impl TailscaleTunnel {
pub fn new(funnel: bool, hostname: Option<String>) -> Self {
Self {
funnel,
hostname,
proc: new_shared_process(),
url: new_shared_url(),
}
}
}
#[async_trait::async_trait]
impl Tunnel for TailscaleTunnel {
fn name(&self) -> &str {
"tailscale"
}
async fn start(&self, local_host: &str, local_port: u16) -> Result<String> {
let subcommand = if self.funnel { "funnel" } else { "serve" };
let hostname = if let Some(ref h) = self.hostname {
h.clone()
} else {
let output = tokio::time::timeout(
tokio::time::Duration::from_secs(10),
Command::new("tailscale")
.args(["status", "--json"])
.output(),
)
.await
.map_err(|_| anyhow::anyhow!("tailscale status --json timed out after 10s"))??;
if !output.status.success() {
bail!(
"tailscale status failed: {}",
String::from_utf8_lossy(&output.stderr)
);
}
let status: serde_json::Value = serde_json::from_slice(&output.stdout)
.map_err(|e| anyhow::anyhow!("Failed to parse tailscale status JSON: {e}"))?;
status["Self"]["DNSName"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("tailscale status missing Self.DNSName field"))?
.trim_end_matches('.')
.to_string()
};
let target = format!("http://{local_host}:{local_port}");
let child = Command::new("tailscale")
.args([subcommand, &target])
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true)
.spawn()?;
let public_url = format!("https://{hostname}");
if let Ok(mut guard) = self.url.write() {
*guard = Some(public_url.clone());
}
let mut guard = self.proc.lock().await;
*guard = Some(TunnelProcess { child });
Ok(public_url)
}
async fn stop(&self) -> Result<()> {
let subcommand = if self.funnel { "funnel" } else { "serve" };
if let Err(e) = Command::new("tailscale")
.args([subcommand, "reset"])
.output()
.await
{
tracing::warn!("tailscale {subcommand} reset failed: {e}");
}
if let Ok(mut guard) = self.url.write() {
*guard = None;
}
kill_shared(&self.proc).await
}
async fn health_check(&self) -> bool {
let guard = self.proc.lock().await;
guard.as_ref().is_some_and(|tp| tp.child.id().is_some())
}
fn public_url(&self) -> Option<String> {
self.url.read().ok().and_then(|guard| guard.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn constructor_stores_hostname_and_mode() {
let tunnel = TailscaleTunnel::new(true, Some("myhost.ts.net".into()));
assert!(tunnel.funnel);
assert_eq!(tunnel.hostname.as_deref(), Some("myhost.ts.net"));
}
#[test]
fn public_url_none_before_start() {
assert!(TailscaleTunnel::new(false, None).public_url().is_none());
}
#[tokio::test]
async fn health_false_before_start() {
assert!(!TailscaleTunnel::new(false, None).health_check().await);
}
#[tokio::test]
async fn stop_without_start_is_ok() {
assert!(TailscaleTunnel::new(false, None).stop().await.is_ok());
}
}
+247
View File
@@ -0,0 +1,247 @@
//! Memory hygiene: automatic cleanup of stale workspace documents.
//!
//! Runs on a configurable cadence and deletes daily log entries older
//! than the retention period. Identity files (`IDENTITY.md`, `SOUL.md`,
//! etc.) are never touched.
//!
//! ```text
//! ┌─────────────────────────────────────────────┐
//! │ Hygiene Pass │
//! │ │
//! │ 1. Check cadence (skip if ran recently) │
//! │ 2. List daily/ documents │
//! │ 3. Delete those older than retention_days │
//! │ 4. Log summary │
//! └─────────────────────────────────────────────┘
//! ```
use std::path::PathBuf;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use crate::workspace::Workspace;
/// Configuration for workspace hygiene.
#[derive(Debug, Clone)]
pub struct HygieneConfig {
/// Whether hygiene is enabled at all.
pub enabled: bool,
/// Documents in `daily/` older than this many days are deleted.
pub retention_days: u32,
/// Minimum hours between hygiene passes.
pub cadence_hours: u32,
/// Directory to store state file (default: `~/.ironclaw`).
pub state_dir: PathBuf,
}
impl Default for HygieneConfig {
fn default() -> Self {
let state_dir = dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw");
Self {
enabled: true,
retention_days: 30,
cadence_hours: 12,
state_dir,
}
}
}
/// Persisted state for tracking hygiene cadence.
#[derive(Debug, Serialize, Deserialize)]
struct HygieneState {
last_run: DateTime<Utc>,
}
/// Summary of what a hygiene pass cleaned up.
#[derive(Debug, Default)]
pub struct HygieneReport {
/// Number of daily log documents deleted.
pub daily_logs_deleted: u32,
/// Whether the run was skipped (cadence not yet elapsed).
pub skipped: bool,
}
impl HygieneReport {
/// True if any cleanup work was done.
pub fn had_work(&self) -> bool {
self.daily_logs_deleted > 0
}
}
/// Run a hygiene pass if the cadence has elapsed.
///
/// This is best-effort: failures are logged but never propagate. The
/// agent should not crash because cleanup failed.
pub async fn run_if_due(workspace: &Workspace, config: &HygieneConfig) -> HygieneReport {
if !config.enabled {
return HygieneReport {
skipped: true,
..Default::default()
};
}
let state_file = config.state_dir.join("memory_hygiene_state.json");
// Check cadence
if let Some(state) = load_state(&state_file) {
let elapsed = Utc::now().signed_duration_since(state.last_run);
let cadence = chrono::Duration::hours(i64::from(config.cadence_hours));
if elapsed < cadence {
tracing::debug!(
hours_since_last = elapsed.num_hours(),
cadence_hours = config.cadence_hours,
"memory hygiene: skipping (cadence not elapsed)"
);
return HygieneReport {
skipped: true,
..Default::default()
};
}
}
tracing::info!(
retention_days = config.retention_days,
"memory hygiene: starting cleanup pass"
);
let mut report = HygieneReport::default();
// Delete old daily logs
match cleanup_daily_logs(workspace, config.retention_days).await {
Ok(count) => report.daily_logs_deleted = count,
Err(e) => tracing::warn!("memory hygiene: failed to clean daily logs: {e}"),
}
if report.had_work() {
tracing::info!(
daily_logs_deleted = report.daily_logs_deleted,
"memory hygiene: cleanup complete"
);
} else {
tracing::debug!("memory hygiene: nothing to clean");
}
// Save state (best-effort)
save_state(&state_file);
report
}
/// Delete daily log documents older than `retention_days`.
async fn cleanup_daily_logs(
workspace: &Workspace,
retention_days: u32,
) -> Result<u32, anyhow::Error> {
let cutoff = Utc::now() - chrono::Duration::days(i64::from(retention_days));
let entries = workspace.list("daily/").await?;
let mut deleted = 0u32;
for entry in entries {
if entry.is_directory {
continue;
}
// Check if the document is old enough to delete
if let Some(updated_at) = entry.updated_at
&& updated_at < cutoff
{
let path = if entry.path.starts_with("daily/") {
entry.path.clone()
} else {
format!("daily/{}", entry.path)
};
if let Err(e) = workspace.delete(&path).await {
tracing::warn!(path, "memory hygiene: failed to delete: {e}");
} else {
tracing::debug!(path, "memory hygiene: deleted old daily log");
deleted += 1;
}
}
}
Ok(deleted)
}
fn state_path_dir(state_file: &std::path::Path) -> Option<&std::path::Path> {
state_file.parent()
}
fn load_state(path: &std::path::Path) -> Option<HygieneState> {
let data = std::fs::read_to_string(path).ok()?;
serde_json::from_str(&data).ok()
}
fn save_state(path: &std::path::Path) {
let state = HygieneState {
last_run: Utc::now(),
};
if let Some(dir) = state_path_dir(path) {
std::fs::create_dir_all(dir).ok();
}
if let Ok(json) = serde_json::to_string_pretty(&state)
&& let Err(e) = std::fs::write(path, json)
{
tracing::warn!("memory hygiene: failed to save state: {e}");
}
}
#[cfg(test)]
mod tests {
use crate::workspace::hygiene::*;
#[test]
fn default_config_is_reasonable() {
let cfg = HygieneConfig::default();
assert!(cfg.enabled);
assert_eq!(cfg.retention_days, 30);
assert_eq!(cfg.cadence_hours, 12);
}
#[test]
fn report_defaults_to_no_work() {
let report = HygieneReport::default();
assert!(!report.had_work());
assert!(!report.skipped);
}
#[test]
fn report_had_work_when_deleted() {
let report = HygieneReport {
daily_logs_deleted: 3,
skipped: false,
};
assert!(report.had_work());
}
#[test]
fn load_state_returns_none_for_missing_file() {
assert!(load_state(std::path::Path::new("/tmp/nonexistent_hygiene.json")).is_none());
}
#[test]
fn save_and_load_state_roundtrip() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("hygiene_state.json");
save_state(&path);
let state = load_state(&path).expect("state should be loadable after save");
// Should be within the last second
let elapsed = Utc::now().signed_duration_since(state.last_run);
assert!(elapsed.num_seconds() < 2);
}
#[test]
fn save_state_creates_parent_dirs() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("nested").join("deep").join("state.json");
save_state(&path);
assert!(path.exists());
}
}
+1
View File
@@ -43,6 +43,7 @@
mod chunker;
mod document;
mod embeddings;
pub mod hygiene;
#[cfg(feature = "postgres")]
mod repository;
mod search;