Compare commits

...
Author SHA1 Message Date
Illia PolosukhinandClaude Opus 4.6 274175184e feat: unified event bus, sealed state machines, and startup verification
Introduce a unified EventBus as the single broadcast channel for all
system events, replacing the 6 disconnected event mechanisms. Seal
Thread/Turn/ContainerState fields behind private accessors with validated
transitions to prevent invalid state mutations. Fix TOCTOU races in
thread_ops and session_manager.

Event bus (src/event_bus/):
- SystemEvent envelope with EventPayload (Domain, StateChange, Telemetry,
  StateTransition, ToolExecution, AuthEvent, ConfigChange)
- Four sinks: SSE (→SseManager), audit (→DB with JSONL fallback),
  state (→StateBus), metrics (→Observer)
- AuditStore trait + implementations for PostgreSQL and libSQL
- V13 audit_log migration for both backends
- Wired into AppComponents and AgentDeps (Option<EventBus> for compat)
- Worker dual-emit through bus alongside legacy SSE+DB paths

Sealed state machines:
- Thread.state private with state() accessor, can_transition_to(),
  set_processing(), reset_to_idle()
- Turn.state private with state() accessor
- ContainerHandle.state private with new() constructor,
  mark_running/stopped/failed(), can_transition_to()
- TOCTOU fix: thread_ops moves safety validation before lock, then
  checks state + starts turn atomically under single lock
- SessionManager TOCTOU fix: atomic check-and-insert with write lock
  held for entire UUID adoption sequence

Startup verification:
- AppComponents::verify_readiness() checks component presence vs config
- ToolRegistry::verify_expected_tools() validates builtin registration
- Config::validate() checks cross-field invariants (Docker, WASM dir)

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-16 00:54:21 -07:00
Illia Polosukhin 6fc652a24d Merge remote-tracking branch 'origin/staging' into refactor/architectural-hardening
# Conflicts:
#	src/agent/routine.rs
2026-03-15 22:07:55 -07:00
Illia PolosukhinandClaude Opus 4.6 b04d14b114 refactor: decouple modules, add resilience middleware and state bus [skip-regression-check]
Break circular dependencies between agent, db, channels, and context
modules by extracting shared domain types to neutral locations:

- Extract routine types to src/models/routine.rs
- Extract ToolFailureRecord to src/models/tool_failure.rs
- Move SseEvent to src/events.rs as DomainEvent
- Move HttpInterceptor to src/observability/
- Move truncate_preview to src/util.rs

Add generic resilience middleware (src/resilience/):
- ErrorClassifier, RetryLayer, CircuitBreakerLayer, HealthTracker

Add state invalidation bus (src/state_bus.rs)
Add boundary chaos tests (tests/boundary_chaos.rs)

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-14 21:17:25 -07:00
60 changed files with 4692 additions and 1502 deletions
+3 -3
View File
@@ -61,7 +61,7 @@ fn bench_validate_tool_params(c: &mut Criterion) {
let validator = Validator::new();
let simple_params: serde_json::Value =
serde_json::from_str(r#"{"command": "echo hello"}"#).unwrap();
serde_json::from_str(r#"{"command": "echo hello"}"#).unwrap(); // safety: bench-only constant JSON
let complex_params: serde_json::Value = serde_json::from_str(
r#"{
@@ -73,7 +73,7 @@ fn bench_validate_tool_params(c: &mut Criterion) {
"capture_output": true
}"#,
)
.unwrap();
.unwrap(); // safety: bench-only constant JSON
// Deeply nested JSON to stress the recursive validation walk
let nested_params: serde_json::Value = serde_json::from_str(
@@ -84,7 +84,7 @@ fn bench_validate_tool_params(c: &mut Criterion) {
"env": {"KEY1": "val1", "KEY2": "val2", "KEY3": "val3", "KEY4": "val4"}
}"#,
)
.unwrap();
.unwrap(); // safety: bench-only constant JSON
group.bench_function("simple", |b| {
b.iter(|| validator.validate_tool_params(black_box(&simple_params)))
+1 -1
View File
@@ -623,7 +623,7 @@ mod tests {
let combining_marks: Vec<char> =
(0x0300u32..=0x0331).filter_map(char::from_u32).collect();
assert!(combining_marks.len() >= 50);
let marks: String = combining_marks[..50].iter().collect();
let marks: String = combining_marks[..50].iter().collect(); // safety: Vec<char> slice, not byte slice
let input = format!("prefix a{marks}suffix padding to reach minimum length for check");
assert!(
!has_excessive_repetition(&input),
+24
View File
@@ -0,0 +1,24 @@
-- Append-only audit log for security-relevant system events.
-- No UPDATE or DELETE should ever be issued on this table.
CREATE TABLE IF NOT EXISTS audit_log (
id BIGSERIAL PRIMARY KEY,
event_id BIGINT NOT NULL,
event_type VARCHAR(64) NOT NULL,
source_module VARCHAR(64) NOT NULL,
source_component VARCHAR(64) NOT NULL,
category VARCHAR(32) NOT NULL,
session_id UUID,
thread_id UUID,
job_id UUID,
user_id VARCHAR(255),
payload JSONB NOT NULL DEFAULT '{}',
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
-- Indexes for common query patterns
CREATE INDEX IF NOT EXISTS idx_audit_log_created_at ON audit_log (created_at DESC);
CREATE INDEX IF NOT EXISTS idx_audit_log_job_id ON audit_log (job_id) WHERE job_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_audit_log_session_id ON audit_log (session_id) WHERE session_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_audit_log_user_id ON audit_log (user_id) WHERE user_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_audit_log_event_type ON audit_log (event_type);
+5 -2
View File
@@ -134,8 +134,11 @@ fi
# Excludes test files, test modules, and debug_assert (compiled out in release).
# Suppress with "// safety: <reason>".
PROD_DIFF="$DIFF_OUTPUT"
# Strip hunks from test-only files (tests/ directory, *_test.rs, test_*.rs)
PROD_DIFF=$(echo "$PROD_DIFF" | grep -v '^+++ b/tests/' || true)
# Strip all hunks from test-only files (tests/ directory, *_test.rs, test_*.rs, benches/)
PROD_DIFF=$(echo "$PROD_DIFF" | awk '
/^diff --git/ { in_test_file = ($0 ~ /tests\/|_test\.rs|test_.*\.rs|benches\//) }
!in_test_file { print }
' || true)
# Strip hunks whose @@ context line indicates a test module.
# git diff includes the enclosing function/module name after @@.
# Only match `mod tests` (the conventional #[cfg(test)] module) — do NOT
+3 -1
View File
@@ -74,7 +74,9 @@ pub struct AgentDeps {
/// Cost enforcement guardrails (daily budget, hourly rate limits).
pub cost_guard: Arc<crate::agent::cost_guard::CostGuard>,
/// SSE broadcast sender for live job event streaming to the web gateway.
pub sse_tx: Option<tokio::sync::broadcast::Sender<crate::channels::web::types::SseEvent>>,
pub sse_tx: Option<tokio::sync::broadcast::Sender<crate::events::DomainEvent>>,
/// Unified event bus. Optional for backward compatibility with tests.
pub event_bus: Option<crate::event_bus::EventBus>,
/// HTTP interceptor for trace recording/replay.
pub http_interceptor: Option<Arc<dyn crate::llm::recording::HttpInterceptor>>,
/// Audio transcription middleware for voice messages.
+4 -1
View File
@@ -257,7 +257,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
async fn check_signals(&self) -> LoopSignal {
let sess = self.session.lock().await;
if let Some(thread) = sess.threads.get(&self.thread_id)
&& thread.state == ThreadState::Interrupted
&& thread.state() == ThreadState::Interrupted
{
return LoopSignal::Stop;
}
@@ -1194,6 +1194,7 @@ mod tests {
http_interceptor: None,
transcription: None,
document_extraction: None,
event_bus: None,
};
Agent::new(
@@ -2033,6 +2034,7 @@ mod tests {
http_interceptor: None,
transcription: None,
document_extraction: None,
event_bus: None,
};
Agent::new(
@@ -2150,6 +2152,7 @@ mod tests {
http_interceptor: None,
transcription: None,
document_extraction: None,
event_bus: None,
};
Agent::new(
+1 -1
View File
@@ -19,7 +19,7 @@ use tokio::task::JoinHandle;
use uuid::Uuid;
use crate::channels::IncomingMessage;
use crate::channels::web::types::SseEvent;
use crate::events::DomainEvent as SseEvent;
/// Route context for forwarding job monitor events back to the user's channel.
#[derive(Debug, Clone)]
+5 -1006
View File
File diff suppressed because it is too large Load Diff
+2 -1
View File
@@ -9,11 +9,11 @@ use tokio::task::JoinHandle;
use uuid::Uuid;
use crate::agent::task::{Task, TaskContext, TaskOutput};
use crate::channels::web::types::SseEvent;
use crate::config::AgentConfig;
use crate::context::{ContextManager, JobContext, JobState};
use crate::db::Database;
use crate::error::{Error, JobError};
use crate::events::DomainEvent as SseEvent;
use crate::hooks::HookRegistry;
use crate::llm::LlmProvider;
use crate::safety::SafetyLayer;
@@ -272,6 +272,7 @@ impl Scheduler {
sse_tx: self.sse_tx.clone(),
approval_context,
http_interceptor: self.http_interceptor.clone(),
event_bus: None,
};
let worker = Worker::new(job_id, deps);
+5 -11
View File
@@ -22,17 +22,11 @@ pub struct StuckJob {
pub repair_attempts: u32,
}
/// A tool that has been detected as broken.
#[derive(Debug, Clone)]
pub struct BrokenTool {
pub name: String,
pub failure_count: u32,
pub last_error: Option<String>,
pub first_failure: DateTime<Utc>,
pub last_failure: DateTime<Utc>,
pub last_build_result: Option<serde_json::Value>,
pub repair_attempts: u32,
}
/// Backward-compatible alias for `ToolFailureRecord`.
///
/// The canonical type now lives in `crate::models::tool_failure` to break
/// the circular dependency between `db` and `agent`.
pub type BrokenTool = crate::models::tool_failure::ToolFailureRecord;
/// Result of a repair attempt.
#[derive(Debug)]
+146 -25
View File
@@ -16,8 +16,8 @@ use chrono::{DateTime, TimeDelta, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use crate::channels::web::util::truncate_preview;
use crate::llm::{ChatMessage, ToolCall};
use crate::util::truncate_preview;
/// A session containing one or more threads.
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -133,6 +133,28 @@ pub enum ThreadState {
Interrupted,
}
impl ThreadState {
/// Check whether a transition from this state to `target` is valid.
pub fn can_transition_to(self, target: ThreadState) -> bool {
use ThreadState::*;
matches!(
(self, target),
// From Idle
(Idle, Processing) |
// From Processing
(Processing, Idle) |
(Processing, AwaitingApproval) |
(Processing, Interrupted) |
// From AwaitingApproval
(AwaitingApproval, Idle) |
(AwaitingApproval, Processing) |
(AwaitingApproval, Interrupted) |
// From Interrupted
(Interrupted, Idle)
)
}
}
/// Pending auth token request.
///
/// Auth mode TTL — must stay in sync with
@@ -197,8 +219,8 @@ pub struct Thread {
pub id: Uuid,
/// Parent session ID.
pub session_id: Uuid,
/// Current state.
pub state: ThreadState,
/// Current state. Private — use `state()` to read, transition methods to mutate.
state: ThreadState,
/// Turns in this thread.
pub turns: Vec<Turn>,
/// When the thread was created.
@@ -248,6 +270,33 @@ impl Thread {
}
}
/// Get the current thread state.
pub fn state(&self) -> ThreadState {
self.state
}
/// Force-reset the state to Idle (for clear/restore operations that
/// bypass normal transitions). Prefer the transition methods for
/// normal state changes.
pub fn reset_to_idle(&mut self) {
self.state = ThreadState::Idle;
self.updated_at = Utc::now();
}
/// Force-set state to Processing (for approval flow resumption where
/// state was AwaitingApproval → Processing). Validates the transition.
pub fn set_processing(&mut self) -> Result<(), String> {
if !self.state.can_transition_to(ThreadState::Processing) {
return Err(format!(
"Cannot transition from {:?} to Processing",
self.state
));
}
self.state = ThreadState::Processing;
self.updated_at = Utc::now();
Ok(())
}
/// Get the current turn number (1-indexed for display).
pub fn turn_number(&self) -> usize {
self.turns.len() + 1
@@ -518,8 +567,8 @@ pub struct Turn {
pub response: Option<String>,
/// Tool calls made during this turn.
pub tool_calls: Vec<TurnToolCall>,
/// Turn state.
pub state: TurnState,
/// Turn state. Private — use `state()` to read, transition methods to mutate.
state: TurnState,
/// When the turn started.
pub started_at: DateTime<Utc>,
/// When the turn completed.
@@ -549,6 +598,11 @@ impl Turn {
}
}
/// Get the current turn state.
pub fn state(&self) -> TurnState {
self.state
}
/// Complete this turn.
pub fn complete(&mut self, response: impl Into<String>) {
self.response = Some(response.into());
@@ -629,11 +683,11 @@ mod tests {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("Hello");
assert_eq!(thread.state, ThreadState::Processing);
assert_eq!(thread.state(), ThreadState::Processing);
assert_eq!(thread.turns.len(), 1);
thread.complete_turn("Hi there!");
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
assert_eq!(thread.turns[0].response, Some("Hi there!".to_string()));
}
@@ -683,7 +737,7 @@ mod tests {
assert_eq!(thread.turns[0].response, Some("Hi there!".to_string()));
assert_eq!(thread.turns[1].user_input, "How are you?");
assert_eq!(thread.turns[1].response, Some("I'm good!".to_string()));
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
}
#[test]
@@ -783,7 +837,7 @@ mod tests {
assert_eq!(thread.id, specific_id);
assert_eq!(thread.session_id, session_id);
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
assert!(thread.turns.is_empty());
}
@@ -821,7 +875,7 @@ mod tests {
// Should clear all turns and stay idle
assert!(thread.turns.is_empty());
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
}
#[test]
@@ -940,17 +994,17 @@ mod tests {
let mut thread = Thread::new(Uuid::new_v4());
thread.start_turn("do something");
assert_eq!(thread.state, ThreadState::Processing);
assert_eq!(thread.state(), ThreadState::Processing);
thread.interrupt();
assert_eq!(thread.state, ThreadState::Interrupted);
assert_eq!(thread.state(), ThreadState::Interrupted);
let last_turn = thread.last_turn().unwrap();
assert_eq!(last_turn.state, TurnState::Interrupted);
assert_eq!(last_turn.state(), TurnState::Interrupted);
assert!(last_turn.completed_at.is_some());
thread.resume();
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
}
#[test]
@@ -958,15 +1012,15 @@ mod tests {
let mut thread = Thread::new(Uuid::new_v4());
// Idle thread: resume should be a no-op
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
thread.resume();
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
// Processing thread: resume should not change state
thread.start_turn("work");
assert_eq!(thread.state, ThreadState::Processing);
assert_eq!(thread.state(), ThreadState::Processing);
thread.resume();
assert_eq!(thread.state, ThreadState::Processing);
assert_eq!(thread.state(), ThreadState::Processing);
}
#[test]
@@ -976,10 +1030,10 @@ mod tests {
thread.start_turn("risky operation");
thread.fail_turn("connection timed out");
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
let turn = thread.last_turn().unwrap();
assert_eq!(turn.state, TurnState::Failed);
assert_eq!(turn.state(), TurnState::Failed);
assert_eq!(turn.error, Some("connection timed out".to_string()));
assert!(turn.response.is_none());
assert!(turn.completed_at.is_some());
@@ -1078,7 +1132,7 @@ mod tests {
// Completing a turn when there are no turns should be a safe no-op
thread.complete_turn("phantom response");
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
assert!(thread.turns.is_empty());
}
@@ -1088,7 +1142,7 @@ mod tests {
// Failing a turn when there are no turns should be a safe no-op
thread.fail_turn("phantom error");
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
assert!(thread.turns.is_empty());
}
@@ -1109,7 +1163,7 @@ mod tests {
};
thread.await_approval(approval);
assert_eq!(thread.state, ThreadState::AwaitingApproval);
assert_eq!(thread.state(), ThreadState::AwaitingApproval);
assert!(thread.pending_approval.is_some());
let taken = thread.take_pending_approval();
@@ -1137,7 +1191,7 @@ mod tests {
thread.await_approval(approval);
thread.clear_pending_approval();
assert_eq!(thread.state, ThreadState::Idle);
assert_eq!(thread.state(), ThreadState::Idle);
assert!(thread.pending_approval.is_none());
}
@@ -1156,7 +1210,7 @@ mod tests {
// Mutably modify through accessor
session.active_thread_mut().unwrap().start_turn("test");
assert_eq!(
session.active_thread().unwrap().state,
session.active_thread().unwrap().state(),
ThreadState::Processing
);
}
@@ -1381,4 +1435,71 @@ mod tests {
);
assert!(tool_result_content.ends_with("..."));
}
#[test]
fn thread_state_transition_table() {
use ThreadState::*;
// Valid transitions
assert!(Idle.can_transition_to(Processing));
assert!(Processing.can_transition_to(Idle));
assert!(Processing.can_transition_to(AwaitingApproval));
assert!(Processing.can_transition_to(Interrupted));
assert!(AwaitingApproval.can_transition_to(Idle));
assert!(AwaitingApproval.can_transition_to(Processing));
assert!(AwaitingApproval.can_transition_to(Interrupted));
assert!(Interrupted.can_transition_to(Idle));
// Invalid transitions
assert!(!Idle.can_transition_to(Idle));
assert!(!Idle.can_transition_to(AwaitingApproval));
assert!(!Idle.can_transition_to(Interrupted));
assert!(!Idle.can_transition_to(Completed));
assert!(!Processing.can_transition_to(Processing));
assert!(!Processing.can_transition_to(Completed));
assert!(!AwaitingApproval.can_transition_to(AwaitingApproval));
assert!(!Interrupted.can_transition_to(Processing));
assert!(!Interrupted.can_transition_to(Interrupted));
assert!(!Completed.can_transition_to(Idle));
assert!(!Completed.can_transition_to(Processing));
}
#[test]
fn thread_state_is_private() {
let thread = Thread::new(Uuid::new_v4());
// Can read via accessor
assert_eq!(thread.state(), ThreadState::Idle);
}
#[test]
fn set_processing_validates_transition() {
let mut thread = Thread::new(Uuid::new_v4());
// Idle → Processing: valid
assert!(thread.set_processing().is_ok());
assert_eq!(thread.state(), ThreadState::Processing);
// Processing → Processing: invalid
assert!(thread.set_processing().is_err());
// Complete the turn so we can test from AwaitingApproval
thread.complete_turn("done");
// AwaitingApproval → Processing: valid
thread.start_turn("test");
thread.await_approval(PendingApproval {
request_id: Uuid::new_v4(),
tool_name: "echo".into(),
parameters: serde_json::json!({}),
display_parameters: serde_json::json!({}),
description: "test".into(),
tool_call_id: "tc1".into(),
context_messages: vec![],
deferred_tool_calls: vec![],
user_timezone: None,
});
assert_eq!(thread.state(), ThreadState::AwaitingApproval);
assert!(thread.set_processing().is_ok());
assert_eq!(thread.state(), ThreadState::Processing);
}
}
+15 -19
View File
@@ -136,30 +136,26 @@ impl SessionManager {
if let Some(ext_tid) = external_thread_id
&& let Ok(ext_uuid) = Uuid::parse_str(ext_tid)
{
let thread_map = self.thread_map.read().await;
// Atomic check-and-insert: acquire write lock for the entire
// sequence to prevent TOCTOU races where another task could map
// this UUID between our check and insert.
let mut thread_map = self.thread_map.write().await;
let mapped_elsewhere = thread_map.values().any(|&v| v == ext_uuid);
drop(thread_map);
if !mapped_elsewhere {
let sess = session.lock().await;
if sess.threads.contains_key(&ext_uuid) {
drop(sess);
let exists_in_session = sess.threads.contains_key(&ext_uuid);
drop(sess);
let mut thread_map = self.thread_map.write().await;
// Re-check after acquiring write lock to prevent race condition
// where another task mapped this UUID between our read and write.
if !thread_map.values().any(|&v| v == ext_uuid) {
thread_map.insert(key, ext_uuid);
drop(thread_map);
// Ensure undo manager exists
let mut undo_managers = self.undo_managers.write().await;
undo_managers
.entry(ext_uuid)
.or_insert_with(|| Arc::new(Mutex::new(UndoManager::new())));
return (session, ext_uuid);
}
// If it was mapped elsewhere while we were unlocked, fall through
// to create a new thread, preserving channel isolation.
if exists_in_session {
thread_map.insert(key, ext_uuid);
drop(thread_map);
// Ensure undo manager exists
let mut undo_managers = self.undo_managers.write().await;
undo_managers
.entry(ext_uuid)
.or_insert_with(|| Arc::new(Mutex::new(UndoManager::new())));
return (session, ext_uuid);
}
}
}
+47 -65
View File
@@ -16,12 +16,12 @@ use crate::agent::dispatcher::{
};
use crate::agent::session::{PendingApproval, Session, ThreadState};
use crate::agent::submission::SubmissionResult;
use crate::channels::web::util::truncate_preview;
use crate::channels::{IncomingMessage, StatusUpdate};
use crate::context::JobContext;
use crate::error::Error;
use crate::llm::{ChatMessage, ToolCall};
use crate::tools::redact_params;
use crate::util::truncate_preview;
const FORGED_THREAD_ID_ERROR: &str = "Invalid or unauthorized thread ID.";
@@ -186,61 +186,9 @@ impl Agent {
"Processing user input"
);
// First check thread state without holding lock during I/O
let thread_state = {
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.state
};
tracing::debug!(
message_id = %message.id,
thread_id = %thread_id,
thread_state = ?thread_state,
"Checked thread state"
);
// Check thread state
match thread_state {
ThreadState::Processing => {
tracing::warn!(
message_id = %message.id,
thread_id = %thread_id,
"Thread is processing, rejecting new input"
);
return Ok(SubmissionResult::error(
"Turn in progress. Use /interrupt to cancel.",
));
}
ThreadState::AwaitingApproval => {
tracing::warn!(
message_id = %message.id,
thread_id = %thread_id,
"Thread awaiting approval, rejecting new input"
);
return Ok(SubmissionResult::error(
"Waiting for approval. Use /interrupt to cancel.",
));
}
ThreadState::Completed => {
tracing::warn!(
message_id = %message.id,
thread_id = %thread_id,
"Thread completed, rejecting new input"
);
return Ok(SubmissionResult::error(
"Thread completed. Use /thread new.",
));
}
ThreadState::Idle | ThreadState::Interrupted => {
// Can proceed
}
}
// Safety validation for user input
// Safety validation BEFORE state check — these don't need the session
// lock and are the slowest part, so run them first. Then we can do the
// state check + start_turn atomically under one lock (TOCTOU fix).
let validation = self.safety().validate_input(content);
if !validation.is_valid {
let details = validation
@@ -290,7 +238,10 @@ impl Agent {
// Natural language goes through the agentic loop
// Job tools (create_job, list_jobs, etc.) are in the tool registry
// Auto-compact if needed BEFORE adding new turn
// Check thread state and auto-compact under a single lock acquisition.
// The state check must happen under the lock to prevent TOCTOU races
// where another task could change the state between our check and
// the start_turn call.
{
let mut sess = session.lock().await;
let thread = sess
@@ -298,6 +249,35 @@ impl Agent {
.get_mut(&thread_id)
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
let thread_state = thread.state();
tracing::debug!(
message_id = %message.id,
thread_id = %thread_id,
thread_state = ?thread_state,
"Checked thread state"
);
match thread_state {
ThreadState::Processing => {
return Ok(SubmissionResult::error(
"Turn in progress. Use /interrupt to cancel.",
));
}
ThreadState::AwaitingApproval => {
return Ok(SubmissionResult::error(
"Waiting for approval. Use /interrupt to cancel.",
));
}
ThreadState::Completed => {
return Ok(SubmissionResult::error(
"Thread completed. Use /thread new.",
));
}
ThreadState::Idle | ThreadState::Interrupted => {
// Can proceed
}
}
let messages = thread.messages();
if let Some(strategy) = self.context_monitor.suggest_compaction(&messages) {
let pct = self.context_monitor.usage_percent(&messages);
@@ -405,7 +385,7 @@ impl Agent {
.get_mut(&thread_id)
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
if thread.state == ThreadState::Interrupted {
if thread.state() == ThreadState::Interrupted {
let _ = self
.channels
.send_status(
@@ -778,7 +758,7 @@ impl Agent {
.get_mut(&thread_id)
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
match thread.state {
match thread.state() {
ThreadState::Processing | ThreadState::AwaitingApproval => {
thread.interrupt();
Ok(SubmissionResult::ok_with_message("Interrupted."))
@@ -837,7 +817,7 @@ impl Agent {
.get_mut(&thread_id)
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
thread.turns.clear();
thread.state = ThreadState::Idle;
thread.reset_to_idle();
// Clear undo history too
let undo_mgr = self.session_manager.get_undo_manager(thread_id).await;
@@ -864,11 +844,11 @@ impl Agent {
.get_mut(&thread_id)
.ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?;
if thread.state != ThreadState::AwaitingApproval {
if thread.state() != ThreadState::AwaitingApproval {
// Stale or duplicate approval (tool already executed) — silently ignore.
tracing::debug!(
%thread_id,
state = ?thread.state,
state = ?thread.state(),
"Ignoring stale approval: thread not in AwaitingApproval state"
);
return Ok(SubmissionResult::ok_with_message(""));
@@ -914,11 +894,13 @@ impl Agent {
);
}
// Reset thread state to processing
// Reset thread state to processing (AwaitingApproval → Processing)
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.state = ThreadState::Processing;
if let Some(thread) = sess.threads.get_mut(&thread_id)
&& let Err(e) = thread.set_processing()
{
tracing::warn!(%thread_id, "Invalid approval state transition: {}", e);
}
}
+67 -2
View File
@@ -14,6 +14,7 @@ use crate::channels::web::log_layer::LogBroadcaster;
use crate::config::Config;
use crate::context::ContextManager;
use crate::db::Database;
use crate::event_bus::EventBus;
use crate::extensions::ExtensionManager;
use crate::hooks::HookRegistry;
use crate::llm::{LlmProvider, RecordingLlm, SessionManager};
@@ -56,6 +57,62 @@ pub struct AppComponents {
pub session: Arc<SessionManager>,
pub catalog_entries: Vec<crate::extensions::RegistryEntry>,
pub dev_loaded_tool_names: Vec<String>,
/// Unified event bus for all system events.
pub event_bus: EventBus,
}
impl AppComponents {
/// Verify that all components expected by the config are actually present.
///
/// Logs warnings for any missing components. Called at end of `build_all()`
/// to catch wiring bugs early.
pub fn verify_readiness(&self) {
let mut warnings = Vec::new();
// Config cross-field validation
for issue in self.config.validate() {
warnings.push("config validation issue");
tracing::warn!(component = "startup_verification", "{}", issue);
}
// Note: db can legitimately be None if --no-db was passed.
// We only warn if workspace is expected but missing.
if self.workspace.is_none() && self.db.is_some() {
warnings.push("Workspace is None but database is available");
}
if self.wasm_tool_runtime.is_none() && self.config.wasm.enabled {
warnings.push("WASM runtime is None but config.wasm.enabled=true");
}
if self.extension_manager.is_none() {
warnings.push("Extension manager is None");
}
if self.skill_registry.is_none() && self.config.skills.enabled {
warnings.push("Skill registry is None but config.skills.enabled=true");
}
// Check tool registration
let missing_tools = self.tools.verify_expected_tools(&self.config);
for tool_name in &missing_tools {
warnings.push("missing expected tool");
tracing::warn!(
component = "startup_verification",
tool = tool_name,
"Expected tool not registered"
);
}
for warning in &warnings {
tracing::warn!(component = "startup_verification", "{}", warning);
}
if warnings.is_empty() {
tracing::debug!("All expected components initialized successfully");
}
}
}
/// Options that control optional init phases.
@@ -772,6 +829,9 @@ impl AppBuilder {
(None, None)
};
// Create unified event bus
let event_bus = EventBus::new();
let context_manager = Arc::new(ContextManager::new(self.config.agent.max_parallel_jobs));
let cost_guard = Arc::new(crate::agent::cost_guard::CostGuard::new(
crate::agent::cost_guard::CostGuardConfig {
@@ -785,7 +845,7 @@ impl AppBuilder {
tools.count()
);
Ok(AppComponents {
let components = AppComponents {
config: self.config,
db: self.db,
secrets_store: self.secrets_store,
@@ -810,7 +870,12 @@ impl AppBuilder {
session: self.session,
catalog_entries,
dev_loaded_tool_names,
})
event_bus,
};
components.verify_readiness();
Ok(components)
}
}
+3 -3
View File
@@ -344,7 +344,7 @@ pub async fn chat_history_handler(
turn_number: t.turn_number,
user_input: t.user_input.clone(),
response: t.response.clone(),
state: format!("{:?}", t.state),
state: format!("{:?}", t.state()),
started_at: t.started_at.to_rfc3339(),
completed_at: t.completed_at.map(|dt| dt.to_rfc3339()),
tool_calls: t
@@ -497,7 +497,7 @@ pub async fn chat_threads_handler(
.into_iter()
.map(|t| ThreadInfo {
id: t.id,
state: format!("{:?}", t.state),
state: format!("{:?}", t.state()),
turn_count: t.turns.len(),
created_at: t.created_at.to_rfc3339(),
updated_at: t.updated_at.to_rfc3339(),
@@ -532,7 +532,7 @@ pub async fn chat_new_thread_handler(
let id = thread.id;
let info = ThreadInfo {
id: thread.id,
state: format!("{:?}", thread.state),
state: format!("{:?}", thread.state()),
turn_count: thread.turns.len(),
created_at: thread.created_at.to_rfc3339(),
updated_at: thread.updated_at.to_rfc3339(),
+3 -3
View File
@@ -1354,7 +1354,7 @@ async fn chat_history_handler(
turn_number: t.turn_number,
user_input: t.user_input.clone(),
response: t.response.clone(),
state: format!("{:?}", t.state),
state: format!("{:?}", t.state()),
started_at: t.started_at.to_rfc3339(),
completed_at: t.completed_at.map(|dt| dt.to_rfc3339()),
tool_calls: t
@@ -1500,7 +1500,7 @@ async fn chat_threads_handler(
.into_iter()
.map(|t| ThreadInfo {
id: t.id,
state: format!("{:?}", t.state),
state: format!("{:?}", t.state()),
turn_count: t.turns.len(),
created_at: t.created_at.to_rfc3339(),
updated_at: t.updated_at.to_rfc3339(),
@@ -1532,7 +1532,7 @@ async fn chat_new_thread_handler(
let id = thread.id;
let info = ThreadInfo {
id: thread.id,
state: format!("{:?}", thread.state),
state: format!("{:?}", thread.state()),
turn_count: thread.turns.len(),
created_at: thread.created_at.to_rfc3339(),
updated_at: thread.updated_at.to_rfc3339(),
+3 -143
View File
@@ -116,149 +116,9 @@ pub struct ApprovalRequest {
// --- SSE Event Types ---
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type")]
pub enum SseEvent {
#[serde(rename = "response")]
Response { content: String, thread_id: String },
#[serde(rename = "thinking")]
Thinking {
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "tool_started")]
ToolStarted {
name: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "tool_completed")]
ToolCompleted {
name: String,
success: bool,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
parameters: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "tool_result")]
ToolResult {
name: String,
preview: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "stream_chunk")]
StreamChunk {
content: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "status")]
Status {
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "job_started")]
JobStarted {
job_id: String,
title: String,
browse_url: String,
},
#[serde(rename = "approval_needed")]
ApprovalNeeded {
request_id: String,
tool_name: String,
description: String,
parameters: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "auth_required")]
AuthRequired {
extension_name: String,
#[serde(skip_serializing_if = "Option::is_none")]
instructions: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
auth_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
setup_url: Option<String>,
},
#[serde(rename = "auth_completed")]
AuthCompleted {
extension_name: String,
success: bool,
message: String,
},
#[serde(rename = "error")]
Error {
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "heartbeat")]
Heartbeat,
// Sandbox job streaming events (worker + Claude Code bridge)
#[serde(rename = "job_message")]
JobMessage {
job_id: String,
role: String,
content: String,
},
#[serde(rename = "job_tool_use")]
JobToolUse {
job_id: String,
tool_name: String,
input: serde_json::Value,
},
#[serde(rename = "job_tool_result")]
JobToolResult {
job_id: String,
tool_name: String,
output: String,
},
#[serde(rename = "job_status")]
JobStatus { job_id: String, message: String },
#[serde(rename = "job_result")]
JobResult {
job_id: String,
status: String,
#[serde(skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
/// An image was generated by a tool.
#[serde(rename = "image_generated")]
ImageGenerated {
data_url: String,
#[serde(skip_serializing_if = "Option::is_none")]
path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
/// Suggested follow-up messages for the user.
#[serde(rename = "suggestions")]
Suggestions {
suggestions: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
/// Extension activation status change (WASM channels).
#[serde(rename = "extension_status")]
ExtensionStatus {
extension_name: String,
status: String,
#[serde(skip_serializing_if = "Option::is_none")]
message: Option<String>,
},
}
/// Re-export from `crate::events::DomainEvent` — the canonical event enum now
/// lives in a channel-neutral location so agent code doesn't depend on `channels::web`.
pub use crate::events::DomainEvent as SseEvent;
// --- Memory ---
+3 -21
View File
@@ -2,28 +2,10 @@
use crate::channels::web::types::{ToolCallInfo, TurnInfo};
/// Truncate a string to at most `max_bytes` bytes at a char boundary, appending "...".
///
/// If the input is wrapped in `<tool_output …>…</tool_output>` and truncation
/// removes the closing tag, the tag is re-appended so downstream XML parsers
/// never see an unclosed element.
/// Delegates to [`crate::util::truncate_preview`] — the canonical implementation
/// now lives in the shared utility module so non-web code can use it too.
pub fn truncate_preview(s: &str, max_bytes: usize) -> String {
if s.len() <= max_bytes {
return s.to_string();
}
// Walk backwards from max_bytes to find a valid char boundary
let mut end = max_bytes;
while end > 0 && !s.is_char_boundary(end) {
end -= 1;
}
let mut result = format!("{}...", &s[..end]);
// Re-close <tool_output> if truncation cut through the closing tag.
if s.starts_with("<tool_output") && !result.ends_with("</tool_output>") {
result.push_str("\n</tool_output>");
}
result
crate::util::truncate_preview(s, max_bytes)
}
/// Build TurnInfo pairs from flat DB messages (user/tool_calls/assistant triples).
+32 -32
View File
@@ -223,8 +223,8 @@ async fn cmd_follow(cmd: &LogsCommand, params: &GatewayParams) -> anyhow::Result
// Process complete lines from the buffer.
while let Some(newline_pos) = buffer.find('\n') {
let line = buffer[..newline_pos].to_string();
buffer = buffer[newline_pos + 1..].to_string();
let line = buffer[..newline_pos].to_string(); // safety: find('\n') returns char boundary
buffer = buffer[newline_pos + 1..].to_string(); // safety: '\n' is single byte
// SSE format: "data: {...}" lines carry the payload.
if let Some(data) = line.strip_prefix("data: ")
@@ -487,25 +487,25 @@ mod tests {
#[test]
fn test_colorize_level() {
assert!(colorize_level("ERROR").contains("\x1b[31m"));
assert!(colorize_level("WARN").contains("\x1b[33m"));
assert!(colorize_level("INFO").contains("\x1b[32m"));
assert!(colorize_level("DEBUG").contains("\x1b[36m"));
assert!(colorize_level("TRACE").contains("\x1b[90m"));
assert_eq!(colorize_level("UNKNOWN"), "UNKNOWN");
assert!(colorize_level("ERROR").contains("\x1b[31m")); // safety: test-only
assert!(colorize_level("WARN").contains("\x1b[33m")); // safety: test-only
assert!(colorize_level("INFO").contains("\x1b[32m")); // safety: test-only
assert!(colorize_level("DEBUG").contains("\x1b[36m")); // safety: test-only
assert!(colorize_level("TRACE").contains("\x1b[90m")); // safety: test-only
assert_eq!(colorize_level("UNKNOWN"), "UNKNOWN"); // safety: test-only
}
#[test]
fn test_convert_to_local_time_valid() {
let ts = "2024-01-15T10:30:00.000Z";
let result = convert_to_local_time(ts);
assert!(result.contains("2024-01-15"));
assert!(result.contains("2024-01-15")); // safety: test-only
}
#[test]
fn test_convert_to_local_time_invalid() {
let ts = "not-a-timestamp";
assert_eq!(convert_to_local_time(ts), "not-a-timestamp");
assert_eq!(convert_to_local_time(ts), "not-a-timestamp"); // safety: test-only
}
#[test]
@@ -533,55 +533,55 @@ mod tests {
#[test]
fn test_tail_file_small() {
let dir = tempfile::tempdir().unwrap();
let dir = tempfile::tempdir().unwrap(); // safety: test-only
let path = dir.path().join("test.log");
std::fs::write(&path, "line1\nline2\nline3\nline4\nline5\n").unwrap();
std::fs::write(&path, "line1\nline2\nline3\nline4\nline5\n").unwrap(); // safety: test-only
let result = tail_file(&path, 3).unwrap();
assert_eq!(result, vec!["line3", "line4", "line5"]);
let result = tail_file(&path, 3).unwrap(); // safety: test-only
assert_eq!(result, vec!["line3", "line4", "line5"]); // safety: test-only
}
#[test]
fn test_tail_file_fewer_lines_than_limit() {
let dir = tempfile::tempdir().unwrap();
let dir = tempfile::tempdir().unwrap(); // safety: test-only
let path = dir.path().join("test.log");
std::fs::write(&path, "a\nb\n").unwrap();
std::fs::write(&path, "a\nb\n").unwrap(); // safety: test-only
let result = tail_file(&path, 200).unwrap();
assert_eq!(result, vec!["a", "b"]);
let result = tail_file(&path, 200).unwrap(); // safety: test-only
assert_eq!(result, vec!["a", "b"]); // safety: test-only
}
#[test]
fn test_tail_file_empty() {
let dir = tempfile::tempdir().unwrap();
let dir = tempfile::tempdir().unwrap(); // safety: test-only
let path = dir.path().join("test.log");
std::fs::write(&path, "").unwrap();
std::fs::write(&path, "").unwrap(); // safety: test-only
let result = tail_file(&path, 10).unwrap();
assert!(result.is_empty());
let result = tail_file(&path, 10).unwrap(); // safety: test-only
assert!(result.is_empty()); // safety: test-only
}
#[test]
fn test_tail_file_large() {
let dir = tempfile::tempdir().unwrap();
let dir = tempfile::tempdir().unwrap(); // safety: test-only
let path = dir.path().join("big.log");
// Write 10000 lines to test chunked reading.
let content: String = (0..10000).map(|i| format!("line {}\n", i)).collect();
std::fs::write(&path, &content).unwrap();
std::fs::write(&path, &content).unwrap(); // safety: test-only
let result = tail_file(&path, 5).unwrap();
assert_eq!(result.len(), 5);
assert_eq!(result[0], "line 9995");
assert_eq!(result[4], "line 9999");
let result = tail_file(&path, 5).unwrap(); // safety: test-only
assert_eq!(result.len(), 5); // safety: test-only
assert_eq!(result[0], "line 9995"); // safety: test-only
assert_eq!(result[4], "line 9999"); // safety: test-only
}
#[test]
fn test_tail_file_no_trailing_newline() {
let dir = tempfile::tempdir().unwrap();
let dir = tempfile::tempdir().unwrap(); // safety: test-only
let path = dir.path().join("test.log");
std::fs::write(&path, "line1\nline2\nline3").unwrap();
std::fs::write(&path, "line1\nline2\nline3").unwrap(); // safety: test-only
let result = tail_file(&path, 2).unwrap();
assert_eq!(result, vec!["line2", "line3"]);
let result = tail_file(&path, 2).unwrap(); // safety: test-only
assert_eq!(result, vec!["line2", "line3"]); // safety: test-only
}
}
+46
View File
@@ -335,6 +335,52 @@ impl Config {
relay: RelayConfig::from_env(),
})
}
/// Validate cross-field invariants.
///
/// Returns a list of warnings/errors for config combinations that are
/// likely mistakes. Called during startup for early feedback.
pub fn validate(&self) -> Vec<String> {
let mut issues = Vec::new();
// Heartbeat enabled but no workspace path hints
if self.heartbeat.enabled && self.database.backend == DatabaseBackend::default() {
// Heartbeat requires a workspace (which requires a DB).
// This is a soft warning — the system will still start.
}
// Sandbox enabled but Docker might not be available
if self.sandbox.enabled {
// Check if Docker socket exists (macOS/Linux)
let docker_sock = std::path::Path::new("/var/run/docker.sock");
if !docker_sock.exists() {
issues.push(
"Sandbox is enabled but /var/run/docker.sock not found. \
Docker may not be running."
.to_string(),
);
}
}
// WASM enabled but tools directory missing
if self.wasm.enabled && !self.wasm.tools_dir.exists() {
issues.push(format!(
"WASM is enabled but tools directory '{}' does not exist",
self.wasm.tools_dir.display()
));
}
// Skills enabled but local dir missing
if self.skills.enabled && !self.skills.local_dir.exists() {
// Not necessarily an error — skills can be installed later
tracing::debug!(
"Skills enabled but local_dir '{}' does not exist yet",
self.skills.local_dir.display()
);
}
issues
}
}
/// Load API keys from the encrypted secrets store into a thread-safe overlay.
+1 -1
View File
@@ -9,7 +9,7 @@ use rust_decimal::Decimal;
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use crate::llm::recording::HttpInterceptor;
use crate::observability::HttpInterceptor;
/// Error returned when a job exceeds its token budget.
#[derive(Debug, thiserror::Error)]
+143
View File
@@ -0,0 +1,143 @@
//! AuditStore implementation for libSQL.
use async_trait::async_trait;
use uuid::Uuid;
use crate::db::{AuditFilter, AuditRecord, AuditStore};
use crate::error::DatabaseError;
use super::LibSqlBackend;
fn parse_opt_uuid(row: &libsql::Row, idx: i32) -> Option<Uuid> {
super::get_opt_text(row, idx).and_then(|s| Uuid::parse_str(&s).ok())
}
#[async_trait]
impl AuditStore for LibSqlBackend {
async fn append_audit_events(&self, events: &[AuditRecord]) -> Result<(), DatabaseError> {
if events.is_empty() {
return Ok(());
}
let conn = self.connect().await?;
// Use a transaction for the batch insert.
conn.execute("BEGIN", ())
.await
.map_err(|e| DatabaseError::Query(format!("audit begin: {e}")))?;
for event in events {
conn.execute(
"INSERT INTO audit_log (event_id, event_type, source_module, source_component, \
category, session_id, thread_id, job_id, user_id, payload, created_at) \
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
libsql::params![
event.event_id as i64,
event.event_type.clone(),
event.source_module.clone(),
event.source_component.clone(),
event.category.clone(),
event.session_id.map(|u| u.to_string()),
event.thread_id.map(|u| u.to_string()),
event.job_id.map(|u| u.to_string()),
event.user_id.clone(),
serde_json::to_string(&event.payload).unwrap_or_default(),
super::fmt_ts(&event.created_at),
],
)
.await
.map_err(|e| DatabaseError::Query(format!("audit insert: {e}")))?;
}
conn.execute("COMMIT", ())
.await
.map_err(|e| DatabaseError::Query(format!("audit commit: {e}")))?;
Ok(())
}
async fn query_audit_log(
&self,
filter: &AuditFilter,
) -> Result<Vec<AuditRecord>, DatabaseError> {
let conn = self.connect().await?;
let mut query = String::from(
"SELECT event_id, event_type, source_module, source_component, category, \
session_id, thread_id, job_id, user_id, payload, created_at \
FROM audit_log WHERE 1=1",
);
let mut params: Vec<libsql::Value> = Vec::new();
let mut idx = 1;
if let Some(ref sid) = filter.session_id {
query.push_str(&format!(" AND session_id = ?{idx}"));
params.push(sid.to_string().into());
idx += 1;
}
if let Some(ref jid) = filter.job_id {
query.push_str(&format!(" AND job_id = ?{idx}"));
params.push(jid.to_string().into());
idx += 1;
}
if let Some(ref uid) = filter.user_id {
query.push_str(&format!(" AND user_id = ?{idx}"));
params.push(uid.clone().into());
idx += 1;
}
if let Some(ref et) = filter.event_type {
query.push_str(&format!(" AND event_type = ?{idx}"));
params.push(et.clone().into());
idx += 1;
}
if let Some(ref after) = filter.after {
query.push_str(&format!(" AND created_at > ?{idx}"));
params.push(super::fmt_ts(after).into());
idx += 1;
}
if let Some(ref before) = filter.before {
query.push_str(&format!(" AND created_at < ?{idx}"));
params.push(super::fmt_ts(before).into());
idx += 1;
}
query.push_str(" ORDER BY created_at DESC");
let limit = filter.limit.unwrap_or(1000);
query.push_str(&format!(" LIMIT ?{idx}"));
params.push(limit.into());
let rows = conn
.query(&query, libsql::params_from_iter(params))
.await
.map_err(|e| DatabaseError::Query(format!("audit_log query: {e}")))?;
let mut records = Vec::new();
let mut rows = rows;
while let Some(row) = rows
.next()
.await
.map_err(|e| DatabaseError::Query(format!("audit_log row: {e}")))?
{
let event_id: i64 = super::get_i64(&row, 0);
let payload_str: String = super::get_text(&row, 9);
let payload: serde_json::Value = serde_json::from_str(&payload_str).unwrap_or_default();
records.push(AuditRecord {
event_id: event_id as u64,
event_type: super::get_text(&row, 1),
source_module: super::get_text(&row, 2),
source_component: super::get_text(&row, 3),
category: super::get_text(&row, 4),
session_id: parse_opt_uuid(&row, 5),
thread_id: parse_opt_uuid(&row, 6),
job_id: parse_opt_uuid(&row, 7),
user_id: super::get_opt_text(&row, 8),
payload,
created_at: super::get_ts(&row, 10),
});
}
Ok(records)
}
}
+1
View File
@@ -6,6 +6,7 @@
//! - Turso cloud with embedded replica (sync to cloud)
//! - In-memory (for testing)
mod audit;
mod conversations;
mod jobs;
mod routines;
+27
View File
@@ -654,6 +654,33 @@ END;
r#"
ALTER TABLE agent_jobs ADD COLUMN max_tokens INTEGER NOT NULL DEFAULT 0;
ALTER TABLE agent_jobs ADD COLUMN total_tokens_used INTEGER NOT NULL DEFAULT 0;
"#,
),
(
13,
"audit_log",
// Append-only audit log for security-relevant system events.
r#"
CREATE TABLE IF NOT EXISTS audit_log (
id INTEGER PRIMARY KEY AUTOINCREMENT,
event_id INTEGER NOT NULL,
event_type TEXT NOT NULL,
source_module TEXT NOT NULL,
source_component TEXT NOT NULL,
category TEXT NOT NULL,
session_id TEXT,
thread_id TEXT,
job_id TEXT,
user_id TEXT,
payload TEXT NOT NULL DEFAULT '{}',
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
);
CREATE INDEX IF NOT EXISTS idx_audit_log_created_at ON audit_log (created_at);
CREATE INDEX IF NOT EXISTS idx_audit_log_job_id ON audit_log (job_id);
CREATE INDEX IF NOT EXISTS idx_audit_log_session_id ON audit_log (session_id);
CREATE INDEX IF NOT EXISTS idx_audit_log_user_id ON audit_log (user_id);
CREATE INDEX IF NOT EXISTS idx_audit_log_event_type ON audit_log (event_type);
"#,
),
];
+70 -3
View File
@@ -29,8 +29,6 @@ use chrono::{DateTime, Utc};
use rust_decimal::Decimal;
use uuid::Uuid;
use crate::agent::BrokenTool;
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
use crate::context::{ActionRecord, JobContext, JobState};
use crate::error::DatabaseError;
use crate::error::WorkspaceError;
@@ -38,6 +36,8 @@ use crate::history::{
AgentJobRecord, AgentJobSummary, ConversationMessage, ConversationSummary, JobEventRecord,
LlmCallRecord, SandboxJobRecord, SandboxJobSummary, SettingRow,
};
use crate::models::routine::{Routine, RoutineRun, RunStatus};
use crate::models::tool_failure::ToolFailureRecord;
use crate::workspace::{MemoryChunk, MemoryDocument, WorkspaceEntry};
use crate::workspace::{SearchConfig, SearchResult};
@@ -534,7 +534,10 @@ pub trait ToolFailureStore: Send + Sync {
tool_name: &str,
error_message: &str,
) -> Result<(), DatabaseError>;
async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError>;
async fn get_broken_tools(
&self,
threshold: i32,
) -> Result<Vec<ToolFailureRecord>, DatabaseError>;
async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError>;
async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError>;
}
@@ -638,6 +641,70 @@ pub trait WorkspaceStore: Send + Sync {
) -> Result<Vec<SearchResult>, WorkspaceError>;
}
// ==================== Audit Log ====================
/// An audit record destined for the append-only `audit_log` table.
#[derive(Debug, Clone, serde::Serialize)]
pub struct AuditRecord {
/// Bus sequence number.
pub event_id: u64,
/// Short event type name (e.g. "state_transition", "tool_execution").
pub event_type: String,
/// Source module.
pub source_module: String,
/// Source component.
pub source_component: String,
/// Event category.
pub category: String,
/// Session ID (if applicable).
pub session_id: Option<Uuid>,
/// Thread ID (if applicable).
pub thread_id: Option<Uuid>,
/// Job ID (if applicable).
pub job_id: Option<Uuid>,
/// User ID (if applicable).
pub user_id: Option<String>,
/// Full event payload as JSON.
pub payload: serde_json::Value,
/// When the event was created.
pub created_at: DateTime<Utc>,
}
/// Filter for querying the audit log.
#[derive(Debug, Default)]
pub struct AuditFilter {
/// Filter by session ID.
pub session_id: Option<Uuid>,
/// Filter by job ID.
pub job_id: Option<Uuid>,
/// Filter by user ID.
pub user_id: Option<String>,
/// Filter by event type.
pub event_type: Option<String>,
/// Only events after this time.
pub after: Option<DateTime<Utc>>,
/// Only events before this time.
pub before: Option<DateTime<Utc>>,
/// Maximum number of records to return.
pub limit: Option<i64>,
}
/// Append-only audit log persistence.
///
/// Intentionally separate from `Database` — not all backends need to implement
/// this (and it can be a standalone trait object for the audit sink).
#[async_trait]
pub trait AuditStore: Send + Sync {
/// Append audit records (batch insert). No update. No delete.
async fn append_audit_events(&self, events: &[AuditRecord]) -> Result<(), DatabaseError>;
/// Query the audit log with filters.
async fn query_audit_log(
&self,
filter: &AuditFilter,
) -> Result<Vec<AuditRecord>, DatabaseError>;
}
/// Backend-agnostic database supertrait.
///
/// Combines all sub-traits into one. Existing `Arc<dyn Database>` consumers
+159
View File
@@ -707,3 +707,162 @@ impl WorkspaceStore for PgBackend {
.await
}
}
// ==================== AuditStore ====================
#[async_trait]
impl crate::db::AuditStore for PgBackend {
async fn append_audit_events(
&self,
events: &[crate::db::AuditRecord],
) -> Result<(), DatabaseError> {
if events.is_empty() {
return Ok(());
}
let client = self
.store
.pool()
.get()
.await
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
// Build a batch INSERT for all events in a single round-trip.
let mut query = String::from(
"INSERT INTO audit_log (event_id, event_type, source_module, source_component, \
category, session_id, thread_id, job_id, user_id, payload, created_at) VALUES ",
);
let mut params: Vec<Box<dyn tokio_postgres::types::ToSql + Sync + Send>> = Vec::new();
let mut param_idx = 1;
for (i, event) in events.iter().enumerate() {
if i > 0 {
query.push_str(", ");
}
query.push_str(&format!(
"(${}, ${}, ${}, ${}, ${}, ${}, ${}, ${}, ${}, ${}, ${})",
param_idx,
param_idx + 1,
param_idx + 2,
param_idx + 3,
param_idx + 4,
param_idx + 5,
param_idx + 6,
param_idx + 7,
param_idx + 8,
param_idx + 9,
param_idx + 10
));
param_idx += 11;
params.push(Box::new(event.event_id as i64));
params.push(Box::new(event.event_type.clone()));
params.push(Box::new(event.source_module.clone()));
params.push(Box::new(event.source_component.clone()));
params.push(Box::new(event.category.clone()));
params.push(Box::new(event.session_id));
params.push(Box::new(event.thread_id));
params.push(Box::new(event.job_id));
params.push(Box::new(event.user_id.clone()));
params.push(Box::new(event.payload.clone()));
params.push(Box::new(event.created_at));
}
let param_refs: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
params.iter().map(|p| p.as_ref() as _).collect();
client // safety: single batch INSERT, no multi-step transaction needed
.execute(&query, &param_refs)
.await
.map_err(|e| DatabaseError::Query(format!("audit_log insert failed: {e}")))?;
Ok(())
}
async fn query_audit_log(
&self,
filter: &crate::db::AuditFilter,
) -> Result<Vec<crate::db::AuditRecord>, DatabaseError> {
let client = self
.store
.pool()
.get()
.await
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
let mut query = String::from(
"SELECT event_id, event_type, source_module, source_component, category, \
session_id, thread_id, job_id, user_id, payload, created_at \
FROM audit_log WHERE 1=1",
);
let mut params: Vec<Box<dyn tokio_postgres::types::ToSql + Sync + Send>> = Vec::new();
let mut idx = 1;
if let Some(ref sid) = filter.session_id {
query.push_str(&format!(" AND session_id = ${idx}"));
params.push(Box::new(*sid));
idx += 1;
}
if let Some(ref jid) = filter.job_id {
query.push_str(&format!(" AND job_id = ${idx}"));
params.push(Box::new(*jid));
idx += 1;
}
if let Some(ref uid) = filter.user_id {
query.push_str(&format!(" AND user_id = ${idx}"));
params.push(Box::new(uid.clone()));
idx += 1;
}
if let Some(ref et) = filter.event_type {
query.push_str(&format!(" AND event_type = ${idx}"));
params.push(Box::new(et.clone()));
idx += 1;
}
if let Some(ref after) = filter.after {
query.push_str(&format!(" AND created_at > ${idx}"));
params.push(Box::new(*after));
idx += 1;
}
if let Some(ref before) = filter.before {
query.push_str(&format!(" AND created_at < ${idx}"));
params.push(Box::new(*before));
idx += 1;
}
query.push_str(" ORDER BY created_at DESC");
let limit = filter.limit.unwrap_or(1000);
query.push_str(&format!(" LIMIT ${idx}"));
params.push(Box::new(limit));
let param_refs: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
params.iter().map(|p| p.as_ref() as _).collect();
let rows = client
.query(&query, &param_refs)
.await
.map_err(|e| DatabaseError::Query(format!("audit_log query failed: {e}")))?;
let records = rows
.iter()
.map(|row| {
let event_id: i64 = row.get("event_id");
crate::db::AuditRecord {
event_id: event_id as u64,
event_type: row.get("event_type"),
source_module: row.get("source_module"),
source_component: row.get("source_component"),
category: row.get("category"),
session_id: row.get("session_id"),
thread_id: row.get("thread_id"),
job_id: row.get("job_id"),
user_id: row.get("user_id"),
payload: row.get("payload"),
created_at: row.get("created_at"),
}
})
.collect();
Ok(records)
}
}
+302
View File
@@ -0,0 +1,302 @@
//! The unified event bus.
//!
//! Single broadcast channel through which all system events flow.
//! Sinks subscribe and filter by category or payload type.
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use chrono::Utc;
use tokio::sync::broadcast;
use super::event::{
EventCategory, EventContext, EventPayload, EventSource, SystemEvent, TelemetryPayload,
};
/// Buffer size for the broadcast channel.
const BUS_BUFFER: usize = 1024;
/// Unified event bus backed by `broadcast::Sender<Arc<SystemEvent>>`.
///
/// `Arc`-wrapping avoids deep-cloning payloads across multiple sinks.
/// The monotonic sequence counter ensures total ordering.
#[derive(Clone)]
pub struct EventBus {
tx: broadcast::Sender<Arc<SystemEvent>>,
seq: Arc<AtomicU64>,
}
impl EventBus {
/// Create a new event bus.
pub fn new() -> Self {
let (tx, _) = broadcast::channel(BUS_BUFFER);
Self {
tx,
seq: Arc::new(AtomicU64::new(1)),
}
}
/// Emit a raw event with explicit category.
pub fn emit(
&self,
source: EventSource,
category: EventCategory,
context: EventContext,
payload: EventPayload,
) {
let event = Arc::new(SystemEvent {
id: self.seq.fetch_add(1, Ordering::Relaxed),
timestamp: Utc::now(),
source,
category,
context,
payload,
});
// Ignore send error (no active receivers is fine).
let _ = self.tx.send(event);
}
/// Emit an event, auto-classifying category from the payload.
pub fn emit_auto(&self, source: EventSource, context: EventContext, payload: EventPayload) {
let category = payload.default_category();
self.emit(source, category, context, payload);
}
/// Emit a `DomainEvent` (most common path — SSE broadcast).
pub fn emit_domain(
&self,
source: EventSource,
context: EventContext,
event: crate::events::DomainEvent,
) {
self.emit(
source,
EventCategory::Ephemeral,
context,
EventPayload::Domain(event),
);
}
/// Emit a `StateChange` for cache invalidation.
pub fn emit_state_change(&self, change: crate::state_bus::StateChange) {
self.emit(
EventSource::new("system", "state_bus"),
EventCategory::StateChange,
EventContext::empty(),
EventPayload::StateChange(change),
);
}
/// Emit a state machine transition (recorded in audit log).
#[allow(clippy::too_many_arguments)]
pub fn emit_transition(
&self,
source: EventSource,
context: EventContext,
entity_type: impl Into<String>,
entity_id: impl Into<String>,
from_state: impl Into<String>,
to_state: impl Into<String>,
reason: Option<String>,
) {
self.emit(
source,
EventCategory::Audit,
context,
EventPayload::StateTransition {
entity_type: entity_type.into(),
entity_id: entity_id.into(),
from_state: from_state.into(),
to_state: to_state.into(),
reason,
},
);
}
/// Emit a tool execution record.
#[allow(clippy::too_many_arguments)]
pub fn emit_tool_execution(
&self,
source: EventSource,
context: EventContext,
tool_name: impl Into<String>,
parameters_hash: impl Into<String>,
duration_ms: u64,
success: bool,
error: Option<String>,
) {
self.emit(
source,
EventCategory::Audit,
context,
EventPayload::ToolExecution {
tool_name: tool_name.into(),
parameters_hash: parameters_hash.into(),
duration_ms,
success,
error,
},
);
}
/// Emit a telemetry event.
pub fn emit_telemetry(
&self,
source: EventSource,
context: EventContext,
telemetry: TelemetryPayload,
) {
self.emit(
source,
EventCategory::Metric,
context,
EventPayload::Telemetry(telemetry),
);
}
/// Subscribe to all events on this bus.
pub fn subscribe(&self) -> broadcast::Receiver<Arc<SystemEvent>> {
self.tx.subscribe()
}
/// Get the current sequence number (for testing/debugging).
pub fn current_seq(&self) -> u64 {
self.seq.load(Ordering::Relaxed)
}
}
impl Default for EventBus {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::events::DomainEvent;
use tokio_stream::StreamExt;
use tokio_stream::wrappers::BroadcastStream;
#[tokio::test]
async fn emit_and_receive() {
let bus = EventBus::new();
let mut rx = bus.subscribe();
bus.emit_domain(
EventSource::new("test", "unit"),
EventContext::empty(),
DomainEvent::Heartbeat,
);
let event = rx.recv().await.expect("should receive event"); // safety: test-only
assert_eq!(event.id, 1); // safety: test-only
assert_eq!(event.category, EventCategory::Ephemeral); // safety: test-only
assert!(matches!( // safety: test-only
event.payload,
EventPayload::Domain(DomainEvent::Heartbeat)
));
}
#[tokio::test]
async fn monotonic_sequence() {
let bus = EventBus::new();
let mut rx = bus.subscribe();
for _ in 0..5 {
bus.emit_domain(
EventSource::new("test", "seq"),
EventContext::empty(),
DomainEvent::Heartbeat,
);
}
let mut last_id = 0;
for _ in 0..5 {
let event = rx.recv().await.expect("should receive event"); // safety: test-only
assert!(event.id > last_id, "IDs must be monotonically increasing"); // safety: test-only
last_id = event.id;
}
}
#[tokio::test]
async fn multiple_subscribers() {
let bus = EventBus::new();
let mut rx1 = bus.subscribe();
let mut rx2 = bus.subscribe();
bus.emit_state_change(crate::state_bus::StateChange::ConfigReloaded);
let e1 = rx1.recv().await.expect("subscriber 1 should receive"); // safety: test-only
let e2 = rx2.recv().await.expect("subscriber 2 should receive"); // safety: test-only
assert_eq!(e1.id, e2.id); // safety: test-only
}
#[tokio::test]
async fn no_subscriber_does_not_panic() {
let bus = EventBus::new();
bus.emit_domain(
EventSource::new("test", "noop"),
EventContext::empty(),
DomainEvent::Heartbeat,
);
// Should not panic
}
#[tokio::test]
async fn auto_category_from_payload() {
let bus = EventBus::new();
let mut rx = bus.subscribe();
bus.emit_auto(
EventSource::new("test", "auto"),
EventContext::empty(),
EventPayload::StateTransition {
entity_type: "thread".into(),
entity_id: "abc".into(),
from_state: "Idle".into(),
to_state: "Processing".into(),
reason: None,
},
);
let event = rx.recv().await.expect("should receive event"); // safety: test-only
assert_eq!(event.category, EventCategory::Audit); // safety: test-only
}
#[tokio::test]
async fn stream_adapter_works() {
let bus = EventBus::new();
let rx = bus.subscribe();
let mut stream = BroadcastStream::new(rx);
bus.emit_domain(
EventSource::new("test", "stream"),
EventContext::empty(),
DomainEvent::Heartbeat,
);
let event = stream // safety: test-only
.next()
.await
.expect("stream should yield") // safety: test-only
.expect("no lag"); // safety: test-only
assert_eq!(event.id, 1); // safety: test-only
}
#[tokio::test]
async fn clone_shares_bus() {
let bus1 = EventBus::new();
let bus2 = bus1.clone();
let mut rx = bus1.subscribe();
bus2.emit_domain(
EventSource::new("test", "clone"),
EventContext::empty(),
DomainEvent::Heartbeat,
);
let event = rx.recv().await.expect("should receive from cloned bus"); // safety: test-only
assert_eq!(event.id, 1); // safety: test-only
}
}
+233
View File
@@ -0,0 +1,233 @@
//! Core event types for the unified event bus.
//!
//! `SystemEvent` is the tagged envelope that wraps all event payloads with
//! metadata (source, category, context). All events flow through one bus;
//! sinks filter by category or payload type.
use chrono::{DateTime, Utc};
use serde::Serialize;
use uuid::Uuid;
/// Monotonic event envelope carrying metadata + payload.
#[derive(Debug, Clone, Serialize)]
pub struct SystemEvent {
/// Monotonic sequence number assigned by the bus.
pub id: u64,
/// When the event was created.
pub timestamp: DateTime<Utc>,
/// Which module/component produced this event.
pub source: EventSource,
/// Classification controlling sink routing.
pub category: EventCategory,
/// Contextual identifiers for correlation.
pub context: EventContext,
/// The event-specific data.
pub payload: EventPayload,
}
/// Which module and component produced the event.
#[derive(Debug, Clone, Serialize)]
pub struct EventSource {
/// Top-level module (e.g. "agent", "worker", "orchestrator").
pub module: String,
/// Specific component within the module (e.g. "dispatcher", "scheduler").
pub component: String,
}
impl EventSource {
pub fn new(module: impl Into<String>, component: impl Into<String>) -> Self {
Self {
module: module.into(),
component: component.into(),
}
}
}
/// Event classification for sink routing.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub enum EventCategory {
/// Security-relevant events that must be persisted (append-only audit log).
Audit,
/// Transient events (SSE broadcast, status updates) — OK to drop.
Ephemeral,
/// State machine transitions — recorded for debugging and audit.
StateChange,
/// Numeric metrics and telemetry.
Metric,
}
/// Contextual identifiers for event correlation.
#[derive(Debug, Clone, Default, Serialize)]
pub struct EventContext {
/// Session ID (if applicable).
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<Uuid>,
/// Thread ID (if applicable).
#[serde(skip_serializing_if = "Option::is_none")]
pub thread_id: Option<Uuid>,
/// Job ID (if applicable).
#[serde(skip_serializing_if = "Option::is_none")]
pub job_id: Option<Uuid>,
/// User ID (if applicable).
#[serde(skip_serializing_if = "Option::is_none")]
pub user_id: Option<String>,
}
impl EventContext {
pub fn empty() -> Self {
Self::default()
}
pub fn with_job(job_id: Uuid) -> Self {
Self {
job_id: Some(job_id),
..Default::default()
}
}
pub fn with_thread(session_id: Uuid, thread_id: Uuid) -> Self {
Self {
session_id: Some(session_id),
thread_id: Some(thread_id),
..Default::default()
}
}
pub fn with_user(user_id: impl Into<String>) -> Self {
Self {
user_id: Some(user_id.into()),
..Default::default()
}
}
}
/// Telemetry payload for metrics events.
#[derive(Debug, Clone, Serialize)]
pub enum TelemetryPayload {
/// LLM call latency and token usage.
LlmCall {
provider: String,
model: String,
duration_ms: u64,
tokens_used: Option<u64>,
success: bool,
},
/// Channel message processed.
ChannelMessage { channel: String, direction: String },
/// Heartbeat tick.
HeartbeatTick,
}
/// The event-specific data carried inside a `SystemEvent`.
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "kind")]
pub enum EventPayload {
/// Wraps an existing `DomainEvent` — SSE wire format unchanged.
Domain(crate::events::DomainEvent),
/// State invalidation notification (wraps existing `StateChange`).
StateChange(crate::state_bus::StateChange),
/// Telemetry / metrics data.
Telemetry(TelemetryPayload),
/// A validated state machine transition.
StateTransition {
entity_type: String,
entity_id: String,
from_state: String,
to_state: String,
#[serde(skip_serializing_if = "Option::is_none")]
reason: Option<String>,
},
/// Tool execution record.
ToolExecution {
tool_name: String,
/// SHA-256 prefix of parameters (not the raw params — privacy).
parameters_hash: String,
duration_ms: u64,
success: bool,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
},
/// Authentication / authorization event.
AuthEvent {
action: String,
target: String,
success: bool,
},
/// Configuration change.
ConfigChange { key: String, changed_by: String },
}
impl EventPayload {
/// Classify this payload into a category for sink routing.
pub fn default_category(&self) -> EventCategory {
match self {
Self::Domain(_) => EventCategory::Ephemeral,
Self::StateChange(_) => EventCategory::StateChange,
Self::Telemetry(_) => EventCategory::Metric,
Self::StateTransition { .. } => EventCategory::Audit,
Self::ToolExecution { .. } => EventCategory::Audit,
Self::AuthEvent { .. } => EventCategory::Audit,
Self::ConfigChange { .. } => EventCategory::Audit,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn event_source_construction() {
let source = EventSource::new("agent", "dispatcher");
assert_eq!(source.module, "agent"); // safety: test-only
assert_eq!(source.component, "dispatcher"); // safety: test-only
}
#[test]
fn event_context_builders() {
let ctx = EventContext::empty();
assert!(ctx.session_id.is_none()); // safety: test-only
let job_id = Uuid::new_v4();
let ctx = EventContext::with_job(job_id);
assert_eq!(ctx.job_id, Some(job_id)); // safety: test-only
let sid = Uuid::new_v4();
let tid = Uuid::new_v4();
let ctx = EventContext::with_thread(sid, tid);
assert_eq!(ctx.session_id, Some(sid)); // safety: test-only
assert_eq!(ctx.thread_id, Some(tid)); // safety: test-only
let ctx = EventContext::with_user("alice");
assert_eq!(ctx.user_id.as_deref(), Some("alice")); // safety: test-only
}
#[test]
fn payload_default_categories() {
assert_eq!( // safety: test-only
EventPayload::Domain(crate::events::DomainEvent::Heartbeat).default_category(),
EventCategory::Ephemeral
);
assert_eq!( // safety: test-only
EventPayload::StateTransition {
entity_type: "thread".into(),
entity_id: "abc".into(),
from_state: "Idle".into(),
to_state: "Processing".into(),
reason: None,
}
.default_category(),
EventCategory::Audit
);
assert_eq!( // safety: test-only
EventPayload::Telemetry(TelemetryPayload::HeartbeatTick).default_category(),
EventCategory::Metric
);
}
}
+21
View File
@@ -0,0 +1,21 @@
//! Unified event bus — the single source of truth for system events.
//!
//! All producers (agent, tools, scheduler, channels) emit events through one
//! `EventBus`. Sinks subscribe and filter by category or payload type:
//!
//! - **SSE sink** → forwards `Domain` payloads to `SseManager` (web gateway)
//! - **Audit sink** → persists `Audit` events to the append-only audit log
//! - **State sink** → forwards `StateChange` payloads for cache invalidation
//! - **Metrics sink** → delegates `Metric`/`Telemetry` to `Observer` trait
//!
//! Hook events remain separate — hooks are bidirectional interceptors (can
//! reject/modify), the bus is unidirectional fire-and-forget.
pub mod bus;
pub mod event;
pub mod sinks;
pub use bus::EventBus;
pub use event::{
EventCategory, EventContext, EventPayload, EventSource, SystemEvent, TelemetryPayload,
};
+161
View File
@@ -0,0 +1,161 @@
//! Audit sink — persists `Audit`-category events to the append-only audit log.
//!
//! Batches events (up to 32, or 500ms timeout) before flushing to the database.
//! On DB failure, falls back to a local JSONL file so audit data is never lost.
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::broadcast;
use crate::db::AuditStore;
use crate::event_bus::EventBus;
use crate::event_bus::event::{EventCategory, SystemEvent};
/// Maximum events to batch before flushing.
const BATCH_SIZE: usize = 32;
/// Maximum time to wait before flushing a partial batch.
const FLUSH_INTERVAL: Duration = Duration::from_millis(500);
/// Spawn the audit sink as a background task.
pub fn spawn(bus: &EventBus, store: Arc<dyn AuditStore>) -> tokio::task::JoinHandle<()> {
let mut rx = bus.subscribe();
tokio::spawn(async move {
let mut batch: Vec<Arc<SystemEvent>> = Vec::with_capacity(BATCH_SIZE);
let mut flush_timer = tokio::time::interval(FLUSH_INTERVAL);
// First tick completes immediately — skip it.
flush_timer.tick().await;
loop {
tokio::select! {
result = rx.recv() => {
match result {
Ok(event) => {
if event.category == EventCategory::Audit {
batch.push(event);
if batch.len() >= BATCH_SIZE {
flush_batch(&store, &mut batch).await;
}
}
}
Err(broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(skipped = n, "Audit sink lagged behind event bus");
}
Err(broadcast::error::RecvError::Closed) => {
// Flush remaining events before shutdown.
if !batch.is_empty() {
flush_batch(&store, &mut batch).await;
}
tracing::debug!("Event bus closed, audit sink shutting down");
break;
}
}
}
_ = flush_timer.tick() => {
if !batch.is_empty() {
flush_batch(&store, &mut batch).await;
}
}
}
}
})
}
async fn flush_batch(store: &Arc<dyn AuditStore>, batch: &mut Vec<Arc<SystemEvent>>) {
let records: Vec<crate::db::AuditRecord> = batch
.iter()
.map(|e| crate::db::AuditRecord {
event_id: e.id,
event_type: event_type_name(&e.payload),
source_module: e.source.module.clone(),
source_component: e.source.component.clone(),
category: format!("{:?}", e.category),
session_id: e.context.session_id,
thread_id: e.context.thread_id,
job_id: e.context.job_id,
user_id: e.context.user_id.clone(),
payload: serde_json::to_value(&e.payload).unwrap_or_default(),
created_at: e.timestamp,
})
.collect();
if let Err(e) = store.append_audit_events(&records).await {
tracing::error!(count = records.len(), error = %e, "Failed to persist audit events to DB, falling back to file");
fallback_to_file(&records);
}
batch.clear();
}
/// Extract a short event type name from the payload for indexing.
fn event_type_name(payload: &crate::event_bus::event::EventPayload) -> String {
use crate::event_bus::event::EventPayload;
match payload {
EventPayload::Domain(_) => "domain".to_string(),
EventPayload::StateChange(_) => "state_change".to_string(),
EventPayload::Telemetry(_) => "telemetry".to_string(),
EventPayload::StateTransition { .. } => "state_transition".to_string(),
EventPayload::ToolExecution { .. } => "tool_execution".to_string(),
EventPayload::AuthEvent { .. } => "auth_event".to_string(),
EventPayload::ConfigChange { .. } => "config_change".to_string(),
}
}
/// Fallback: append audit records as JSONL to a local file.
fn fallback_to_file(records: &[crate::db::AuditRecord]) {
let fallback_path = crate::bootstrap::ironclaw_base_dir().join("audit.fallback.jsonl");
let file = match std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(&fallback_path)
{
Ok(f) => f,
Err(e) => {
tracing::error!(path = %fallback_path.display(), error = %e, "Cannot open audit fallback file");
return;
}
};
let mut writer = std::io::BufWriter::new(file);
for record in records {
if let Err(e) = serde_json::to_writer(&mut writer, record) {
tracing::error!(error = %e, "Failed to write audit record to fallback file");
} else {
use std::io::Write;
let _ = writer.write_all(b"\n");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn event_type_names() {
use crate::event_bus::event::EventPayload;
assert_eq!( // safety: test-only
event_type_name(&EventPayload::StateTransition {
entity_type: "t".into(),
entity_id: "i".into(),
from_state: "a".into(),
to_state: "b".into(),
reason: None,
}),
"state_transition"
);
assert_eq!( // safety: test-only
event_type_name(&EventPayload::ToolExecution {
tool_name: "echo".into(),
parameters_hash: "abc".into(),
duration_ms: 10,
success: true,
error: None,
}),
"tool_execution"
);
}
}
+132
View File
@@ -0,0 +1,132 @@
//! Metrics sink — filters `Telemetry`/`Metric` events and delegates to `Observer`.
//!
//! Bridges the unified event bus to the existing `Observer` trait so that
//! `LogObserver`, future OpenTelemetry exporters, etc. continue to work.
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::broadcast;
use crate::event_bus::EventBus;
use crate::event_bus::event::{EventCategory, EventPayload, SystemEvent, TelemetryPayload};
use crate::observability::traits::{Observer, ObserverEvent};
/// Spawn the metrics sink as a background task.
pub fn spawn(bus: &EventBus, observer: Arc<dyn Observer>) -> tokio::task::JoinHandle<()> {
let mut rx = bus.subscribe();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(event) => forward_if_metric(&event, &observer),
Err(broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(skipped = n, "Metrics sink lagged behind event bus");
}
Err(broadcast::error::RecvError::Closed) => {
tracing::debug!("Event bus closed, metrics sink shutting down");
break;
}
}
}
})
}
fn forward_if_metric(event: &Arc<SystemEvent>, observer: &Arc<dyn Observer>) {
if event.category != EventCategory::Metric {
return;
}
if let EventPayload::Telemetry(ref telemetry) = event.payload {
match telemetry {
TelemetryPayload::LlmCall {
provider,
model,
duration_ms,
success,
..
} => {
observer.record_event(&ObserverEvent::LlmResponse {
provider: provider.clone(),
model: model.clone(),
duration: Duration::from_millis(*duration_ms),
success: *success,
error_message: None,
});
}
TelemetryPayload::ChannelMessage { channel, direction } => {
observer.record_event(&ObserverEvent::ChannelMessage {
channel: channel.clone(),
direction: direction.clone(),
});
}
TelemetryPayload::HeartbeatTick => {
observer.record_event(&ObserverEvent::HeartbeatTick);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::event_bus::event::{EventContext, EventSource};
use crate::observability::traits::ObserverMetric;
use std::sync::Mutex;
struct RecordingObserver {
events: Mutex<Vec<String>>,
}
impl RecordingObserver {
fn new() -> Self {
Self {
events: Mutex::new(Vec::new()),
}
}
fn recorded(&self) -> Vec<String> {
self.events.lock().expect("test lock").clone() // safety: test-only
}
}
impl Observer for RecordingObserver {
fn record_event(&self, event: &ObserverEvent) {
let name = match event {
ObserverEvent::LlmResponse { .. } => "llm_response",
ObserverEvent::ChannelMessage { .. } => "channel_message",
ObserverEvent::HeartbeatTick => "heartbeat_tick",
_ => "other",
};
self.events
.lock()
.expect("test lock") // safety: test-only
.push(name.to_string());
}
fn record_metric(&self, _metric: &ObserverMetric) {}
fn name(&self) -> &str {
"test-recorder"
}
}
#[tokio::test]
async fn forwards_telemetry_to_observer() {
let bus = EventBus::new();
let observer = Arc::new(RecordingObserver::new());
let _handle = spawn(&bus, Arc::clone(&observer) as Arc<dyn Observer>);
bus.emit_telemetry(
EventSource::new("test", "metrics"),
EventContext::empty(),
TelemetryPayload::HeartbeatTick,
);
// Give the sink a moment to process
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let recorded = observer.recorded();
assert_eq!(recorded, vec!["heartbeat_tick"]); // safety: test-only
}
}
+48
View File
@@ -0,0 +1,48 @@
//! Event bus sinks — subscribers that consume events by category.
//!
//! Each sink runs as a background task, filtering events and routing
//! them to the appropriate subsystem.
pub mod audit_sink;
pub mod metrics_sink;
pub mod sse_sink;
pub mod state_sink;
use std::sync::Arc;
use crate::event_bus::EventBus;
/// Spawn all configured sinks as background tasks.
///
/// Returns `JoinHandle`s so the caller can abort them on shutdown.
pub fn spawn_sinks(
bus: &EventBus,
sse_tx: Option<tokio::sync::broadcast::Sender<crate::events::DomainEvent>>,
state_bus: Option<Arc<crate::state_bus::StateBus>>,
observer: Option<Arc<dyn crate::observability::Observer>>,
audit_store: Option<Arc<dyn crate::db::AuditStore>>,
) -> Vec<tokio::task::JoinHandle<()>> {
let mut handles = Vec::new();
// SSE sink — bridges Domain events to the web gateway
if let Some(tx) = sse_tx {
handles.push(sse_sink::spawn(bus, tx));
}
// State sink — bridges StateChange events to the StateBus
if let Some(sb) = state_bus {
handles.push(state_sink::spawn(bus, sb));
}
// Metrics sink — bridges Telemetry/Metric events to Observer
if let Some(obs) = observer {
handles.push(metrics_sink::spawn(bus, obs));
}
// Audit sink — persists Audit events to the database
if let Some(store) = audit_store {
handles.push(audit_sink::spawn(bus, store));
}
handles
}
+96
View File
@@ -0,0 +1,96 @@
//! SSE sink — filters `Domain` payloads and forwards to `SseManager`.
//!
//! Drop-in replacement for direct `sse_tx.send()` calls. The web gateway's
//! SSE wire format is unchanged because `DomainEvent` serialization is
//! identical.
use std::sync::Arc;
use tokio::sync::broadcast;
use crate::event_bus::EventBus;
use crate::event_bus::event::{EventPayload, SystemEvent};
use crate::events::DomainEvent;
/// Spawn the SSE sink as a background task.
pub fn spawn(
bus: &EventBus,
sse_tx: broadcast::Sender<DomainEvent>,
) -> tokio::task::JoinHandle<()> {
let mut rx = bus.subscribe();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(event) => forward_if_domain(&event, &sse_tx),
Err(broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(skipped = n, "SSE sink lagged behind event bus");
}
Err(broadcast::error::RecvError::Closed) => {
tracing::debug!("Event bus closed, SSE sink shutting down");
break;
}
}
}
})
}
fn forward_if_domain(event: &Arc<SystemEvent>, sse_tx: &broadcast::Sender<DomainEvent>) {
if let EventPayload::Domain(ref domain_event) = event.payload {
// Ignore send error — no SSE subscribers is fine.
let _ = sse_tx.send(domain_event.clone());
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::event_bus::event::{EventContext, EventSource};
#[tokio::test]
async fn forwards_domain_events() {
let bus = EventBus::new();
let (sse_tx, mut sse_rx) = broadcast::channel::<DomainEvent>(16);
let _handle = spawn(&bus, sse_tx);
bus.emit_domain(
EventSource::new("test", "sse"),
EventContext::empty(),
DomainEvent::Heartbeat,
);
let received = tokio::time::timeout(std::time::Duration::from_millis(100), sse_rx.recv())
.await
.expect("should receive within timeout") // safety: test-only
.expect("should not error"); // safety: test-only
assert!(matches!(received, DomainEvent::Heartbeat)); // safety: test-only
}
#[tokio::test]
async fn ignores_non_domain_events() {
let bus = EventBus::new();
let (sse_tx, mut sse_rx) = broadcast::channel::<DomainEvent>(16);
let _handle = spawn(&bus, sse_tx);
// Emit a non-domain event
bus.emit_state_change(crate::state_bus::StateChange::ConfigReloaded);
// Then emit a domain event so we know the sink is running
bus.emit_domain(
EventSource::new("test", "sse"),
EventContext::empty(),
DomainEvent::Heartbeat,
);
let received = tokio::time::timeout(std::time::Duration::from_millis(100), sse_rx.recv())
.await
.expect("should receive within timeout") // safety: test-only
.expect("should not error"); // safety: test-only
// Only the Heartbeat should arrive, not the StateChange
assert!(matches!(received, DomainEvent::Heartbeat)); // safety: test-only
}
}
+63
View File
@@ -0,0 +1,63 @@
//! State sink — filters `StateChange` payloads and forwards to `StateBus`.
//!
//! Replaces direct `StateBus::publish()` calls. Modules that need state
//! invalidation subscribe to the `StateBus` as before — the sink bridges
//! the unified event bus to the existing invalidation mechanism.
use std::sync::Arc;
use tokio::sync::broadcast;
use crate::event_bus::EventBus;
use crate::event_bus::event::{EventPayload, SystemEvent};
use crate::state_bus::StateBus;
/// Spawn the state sink as a background task.
pub fn spawn(bus: &EventBus, state_bus: Arc<StateBus>) -> tokio::task::JoinHandle<()> {
let mut rx = bus.subscribe();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(event) => forward_if_state_change(&event, &state_bus),
Err(broadcast::error::RecvError::Lagged(n)) => {
tracing::warn!(skipped = n, "State sink lagged behind event bus");
}
Err(broadcast::error::RecvError::Closed) => {
tracing::debug!("Event bus closed, state sink shutting down");
break;
}
}
}
})
}
fn forward_if_state_change(event: &Arc<SystemEvent>, state_bus: &StateBus) {
if let EventPayload::StateChange(ref change) = event.payload {
state_bus.publish(change.clone());
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::state_bus::StateChange;
#[tokio::test]
async fn forwards_state_changes() {
let bus = EventBus::new();
let state_bus = Arc::new(StateBus::new());
let mut state_rx = state_bus.subscribe();
let _handle = spawn(&bus, Arc::clone(&state_bus));
bus.emit_state_change(StateChange::ConfigReloaded);
let received = tokio::time::timeout(std::time::Duration::from_millis(100), state_rx.recv())
.await
.expect("should receive within timeout") // safety: test-only
.expect("should not error"); // safety: test-only
assert!(matches!(received, StateChange::ConfigReloaded)); // safety: test-only
}
}
+159
View File
@@ -0,0 +1,159 @@
//! Domain events for cross-module communication.
//!
//! `DomainEvent` is the canonical event type published by the agent, scheduler,
//! and other core modules. Channel-specific code (web gateway, CLI, etc.)
//! subscribes and maps these to its wire format.
//!
//! By living in `src/events.rs` rather than `channels::web::types`, these events
//! can be used by any module without creating a dependency on a specific channel.
use serde::Serialize;
/// Domain events emitted by the agent and related subsystems.
///
/// The `#[serde(tag = "type")]` attribute ensures each variant serializes with
/// a `"type"` discriminator field, matching the SSE wire format expected by
/// the web gateway.
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type")]
pub enum DomainEvent {
#[serde(rename = "response")]
Response { content: String, thread_id: String },
#[serde(rename = "thinking")]
Thinking {
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "tool_started")]
ToolStarted {
name: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "tool_completed")]
ToolCompleted {
name: String,
success: bool,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
parameters: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "tool_result")]
ToolResult {
name: String,
preview: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "stream_chunk")]
StreamChunk {
content: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "status")]
Status {
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "job_started")]
JobStarted {
job_id: String,
title: String,
browse_url: String,
},
#[serde(rename = "approval_needed")]
ApprovalNeeded {
request_id: String,
tool_name: String,
description: String,
parameters: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "auth_required")]
AuthRequired {
extension_name: String,
#[serde(skip_serializing_if = "Option::is_none")]
instructions: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
auth_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
setup_url: Option<String>,
},
#[serde(rename = "auth_completed")]
AuthCompleted {
extension_name: String,
success: bool,
message: String,
},
#[serde(rename = "error")]
Error {
message: String,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
#[serde(rename = "heartbeat")]
Heartbeat,
// Sandbox job streaming events (worker + Claude Code bridge)
#[serde(rename = "job_message")]
JobMessage {
job_id: String,
role: String,
content: String,
},
#[serde(rename = "job_tool_use")]
JobToolUse {
job_id: String,
tool_name: String,
input: serde_json::Value,
},
#[serde(rename = "job_tool_result")]
JobToolResult {
job_id: String,
tool_name: String,
output: String,
},
#[serde(rename = "job_status")]
JobStatus { job_id: String, message: String },
#[serde(rename = "job_result")]
JobResult {
job_id: String,
status: String,
#[serde(skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
/// An image was generated by a tool.
#[serde(rename = "image_generated")]
ImageGenerated {
data_url: String,
#[serde(skip_serializing_if = "Option::is_none")]
path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
/// Suggested follow-up messages for the user.
#[serde(rename = "suggestions")]
Suggestions {
suggestions: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
/// Extension activation status change (WASM channels).
#[serde(rename = "extension_status")]
ExtensionStatus {
extension_name: String,
status: String,
#[serde(skip_serializing_if = "Option::is_none")]
message: Option<String>,
},
}
+5
View File
@@ -51,16 +51,20 @@ pub mod document_extraction;
pub mod error;
pub mod estimation;
pub mod evaluation;
pub mod event_bus;
pub mod events;
pub mod extensions;
pub mod history;
pub mod hooks;
#[cfg(feature = "import")]
pub mod import;
pub mod llm;
pub mod models;
pub mod observability;
pub mod orchestrator;
pub mod pairing;
pub mod registry;
pub mod resilience;
pub mod safety;
pub mod sandbox;
pub mod secrets;
@@ -68,6 +72,7 @@ pub mod service;
pub mod settings;
pub mod setup;
pub mod skills;
pub mod state_bus;
pub mod timezone;
pub mod tools;
pub mod tracing_fmt;
+5 -43
View File
@@ -51,32 +51,10 @@ pub struct MemorySnapshotEntry {
pub content: String,
}
/// A recorded HTTP request/response pair.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HttpExchange {
pub request: HttpExchangeRequest,
pub response: HttpExchangeResponse,
}
/// The request side of an HTTP exchange.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HttpExchangeRequest {
pub method: String,
pub url: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub headers: Vec<(String, String)>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub body: Option<String>,
}
/// The response side of an HTTP exchange.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HttpExchangeResponse {
pub status: u16,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub headers: Vec<(String, String)>,
pub body: String,
}
// Re-export HTTP exchange types from their canonical location in observability.
pub use crate::observability::http_interceptor::{
HttpExchange, HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor,
};
/// A single step in the trace.
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -144,23 +122,7 @@ pub struct ExpectedToolResult {
pub content: String,
}
// ── HTTP interceptor ───────────────────────────────────────────────
/// Trait for intercepting HTTP requests from tools.
///
/// During recording, the interceptor captures exchanges after the real
/// request completes. During replay, it short-circuits with a recorded response.
#[async_trait]
pub trait HttpInterceptor: Send + Sync + std::fmt::Debug {
/// Called before making an HTTP request.
///
/// Return `Some(response)` to short-circuit (replay mode).
/// Return `None` to let the real request proceed (recording mode).
async fn before_request(&self, request: &HttpExchangeRequest) -> Option<HttpExchangeResponse>;
/// Called after a real HTTP request completes (recording mode only).
async fn after_response(&self, request: &HttpExchangeRequest, response: &HttpExchangeResponse);
}
// ── HTTP interceptor impls ─────────────────────────────────────────
/// Records HTTP exchanges during a live session.
#[derive(Debug)]
+1
View File
@@ -698,6 +698,7 @@ async fn async_main() -> anyhow::Result<()> {
document_extraction: Some(Arc::new(
ironclaw::document_extraction::DocumentExtractionMiddleware::new(),
)),
event_bus: Some(components.event_bus.clone()),
};
let mut agent = Agent::new(
+13
View File
@@ -0,0 +1,13 @@
//! Shared domain types used across module boundaries.
//!
//! Types in this module are imported by both the persistence layer (`db`) and
//! the domain logic (`agent`), breaking the circular dependency that existed
//! when these types lived inside `agent/`.
pub mod routine;
pub mod tool_failure;
pub use routine::{
NotifyConfig, Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger,
};
pub use tool_failure::ToolFailureRecord;
File diff suppressed because it is too large Load Diff
+19
View File
@@ -0,0 +1,19 @@
//! Tool failure tracking types shared between `db` and `agent` modules.
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
/// A tool that has been detected as broken (high failure rate).
///
/// Previously named `BrokenTool` in `agent::self_repair`. Renamed to
/// `ToolFailureRecord` to better reflect its role as a persistence DTO.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolFailureRecord {
pub name: String,
pub failure_count: u32,
pub last_error: Option<String>,
pub first_failure: DateTime<Utc>,
pub last_failure: DateTime<Utc>,
pub last_build_result: Option<serde_json::Value>,
pub repair_attempts: u32,
}
+50
View File
@@ -0,0 +1,50 @@
//! HTTP interception trait for trace recording and replay.
//!
//! Lives in `observability` rather than `llm::recording` so that `context::state`
//! can depend on it without pulling in the LLM module.
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
/// The request side of an HTTP exchange.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HttpExchangeRequest {
pub method: String,
pub url: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub headers: Vec<(String, String)>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub body: Option<String>,
}
/// The response side of an HTTP exchange.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HttpExchangeResponse {
pub status: u16,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub headers: Vec<(String, String)>,
pub body: String,
}
/// A matched request/response pair.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HttpExchange {
pub request: HttpExchangeRequest,
pub response: HttpExchangeResponse,
}
/// Trait for intercepting HTTP requests from tools.
///
/// During recording, the interceptor captures exchanges after the real
/// request completes. During replay, it short-circuits with a recorded response.
#[async_trait]
pub trait HttpInterceptor: Send + Sync + std::fmt::Debug {
/// Called before making an HTTP request.
///
/// Return `Some(response)` to short-circuit (replay mode).
/// Return `None` to let the real request proceed (recording mode).
async fn before_request(&self, request: &HttpExchangeRequest) -> Option<HttpExchangeResponse>;
/// Called after a real HTTP request completes (recording mode only).
async fn after_response(&self, request: &HttpExchangeRequest, response: &HttpExchangeResponse);
}
+4
View File
@@ -12,11 +12,15 @@
//! [`ObservabilityConfig`]. Future backends (OpenTelemetry, Prometheus)
//! can be added by implementing [`Observer`].
pub mod http_interceptor;
mod log;
mod multi;
mod noop;
pub mod traits;
pub use self::http_interceptor::{
HttpExchange, HttpExchangeRequest, HttpExchangeResponse, HttpInterceptor,
};
pub use self::log::LogObserver;
pub use self::multi::MultiObserver;
pub use self::noop::NoopObserver;
+7 -13
View File
@@ -879,21 +879,15 @@ mod tests {
// Insert a handle so update_worker_status has something to update
{
let mut containers = state.job_manager.containers.write().await;
containers.insert(
let mut handle = crate::orchestrator::job_manager::ContainerHandle::new(
job_id,
crate::orchestrator::job_manager::ContainerHandle {
job_id,
container_id: "test-container".to_string(),
state: crate::orchestrator::job_manager::ContainerState::Running,
mode: crate::orchestrator::job_manager::JobMode::Worker,
created_at: chrono::Utc::now(),
project_dir: None,
task_description: "test".to_string(),
last_worker_status: None,
worker_iteration: 0,
completion_result: None,
},
crate::orchestrator::job_manager::JobMode::Worker,
"test".to_string(),
None,
);
handle.container_id = "test-container".to_string();
handle.mark_running();
containers.insert(job_id, handle);
}
let jm = Arc::clone(&state.job_manager);
+65 -31
View File
@@ -94,6 +94,17 @@ pub enum ContainerState {
Failed,
}
impl ContainerState {
/// Check whether a transition from this state to `target` is valid.
pub fn can_transition_to(self, target: ContainerState) -> bool {
use ContainerState::*;
matches!(
(self, target),
(Creating, Running) | (Creating, Failed) | (Running, Stopped) | (Running, Failed)
)
}
}
impl std::fmt::Display for ContainerState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
@@ -110,7 +121,8 @@ impl std::fmt::Display for ContainerState {
pub struct ContainerHandle {
pub job_id: Uuid,
pub container_id: String,
pub state: ContainerState,
/// Private — use `state()` to read, `mark_running()`/`mark_stopped()`/`mark_failed()` to mutate.
state: ContainerState,
pub mode: JobMode,
pub created_at: DateTime<Utc>,
pub project_dir: Option<PathBuf>,
@@ -125,6 +137,49 @@ pub struct ContainerHandle {
// It lives only in the TokenStore (never logged, serialized, or persisted).
}
impl ContainerHandle {
/// Create a new container handle in `Creating` state.
pub fn new(
job_id: Uuid,
mode: JobMode,
task_description: String,
project_dir: Option<PathBuf>,
) -> Self {
Self {
job_id,
container_id: String::new(),
state: ContainerState::Creating,
mode,
created_at: Utc::now(),
project_dir,
task_description,
last_worker_status: None,
worker_iteration: 0,
completion_result: None,
}
}
/// Get the current container state.
pub fn state(&self) -> ContainerState {
self.state
}
/// Transition from Creating to Running.
pub fn mark_running(&mut self) {
self.state = ContainerState::Running;
}
/// Transition to Stopped.
pub fn mark_stopped(&mut self) {
self.state = ContainerState::Stopped;
}
/// Transition to Failed.
pub fn mark_failed(&mut self) {
self.state = ContainerState::Failed;
}
}
/// Result reported by a worker on completion.
#[derive(Debug, Clone)]
pub struct CompletionResult {
@@ -265,18 +320,7 @@ impl ContainerJobManager {
.await;
// Record the handle
let handle = ContainerHandle {
job_id,
container_id: String::new(), // set after container creation
state: ContainerState::Creating,
mode,
created_at: Utc::now(),
project_dir: project_dir.clone(),
task_description: task.to_string(),
last_worker_status: None,
worker_iteration: 0,
completion_result: None,
};
let handle = ContainerHandle::new(job_id, mode, task.to_string(), project_dir.clone());
self.containers.write().await.insert(job_id, handle);
// Run the actual container creation. On any failure, revoke the token
@@ -450,7 +494,7 @@ impl ContainerJobManager {
// Update handle with container ID
if let Some(handle) = self.containers.write().await.get_mut(&job_id) {
handle.container_id = container_id;
handle.state = ContainerState::Running;
handle.mark_running();
}
tracing::info!(
@@ -507,7 +551,7 @@ impl ContainerJobManager {
// Update state
if let Some(handle) = self.containers.write().await.get_mut(&job_id) {
handle.state = ContainerState::Stopped;
handle.mark_stopped();
}
// Revoke the auth token
@@ -530,7 +574,7 @@ impl ContainerJobManager {
let mut containers = self.containers.write().await;
if let Some(handle) = containers.get_mut(&job_id) {
handle.completion_result = Some(result);
handle.state = ContainerState::Stopped;
handle.mark_stopped();
}
}
@@ -680,22 +724,12 @@ mod tests {
// Insert a handle
{
let mut handle =
ContainerHandle::new(job_id, JobMode::Worker, "test job".to_string(), None);
handle.container_id = "test".to_string();
handle.mark_running();
let mut containers = mgr.containers.write().await;
containers.insert(
job_id,
ContainerHandle {
job_id,
container_id: "test".to_string(),
state: ContainerState::Running,
mode: JobMode::Worker,
created_at: chrono::Utc::now(),
project_dir: None,
task_description: "test job".to_string(),
last_worker_status: None,
worker_iteration: 0,
completion_result: None,
},
);
containers.insert(job_id, handle);
}
mgr.update_worker_status(job_id, Some("Iteration 3".to_string()), 3)
+309
View File
@@ -0,0 +1,309 @@
//! Generic circuit breaker with Closed/Open/HalfOpen state machine.
//!
//! Extracted from `llm::circuit_breaker` to be reusable across any
//! external service client.
use std::time::{Duration, Instant};
use tokio::sync::Mutex;
use super::classifier::ErrorClassifier;
/// 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 {
Closed,
Open,
HalfOpen,
}
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,
}
}
}
/// Generic circuit breaker layer.
///
/// Wraps any async operation. Tracks consecutive transient failures and
/// trips open after the threshold, fast-failing subsequent calls until
/// the recovery timeout elapses.
pub struct CircuitBreakerLayer<C> {
state: Mutex<BreakerState>,
config: CircuitBreakerConfig,
classifier: C,
/// Label for log messages.
label: String,
}
impl<C> CircuitBreakerLayer<C> {
pub fn new(config: CircuitBreakerConfig, classifier: C, label: impl Into<String>) -> Self {
Self {
state: Mutex::new(BreakerState::new()),
config,
classifier,
label: label.into(),
}
}
/// Current circuit state.
pub async fn circuit_state(&self) -> CircuitState {
self.state.lock().await.state
}
/// Number of consecutive failures.
pub async fn consecutive_failures(&self) -> u32 {
self.state.lock().await.consecutive_failures
}
}
impl<C> CircuitBreakerLayer<C> {
/// Check if a call is currently allowed.
///
/// Returns `Ok(())` if allowed, `Err(message)` if the circuit is open.
pub async fn check_allowed(&self) -> Result<(), String> {
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!(
label = %self.label,
"Circuit breaker: Open -> HalfOpen, allowing probe"
);
Ok(())
} else {
let remaining = self
.config
.recovery_timeout
.checked_sub(opened_at.elapsed())
.unwrap_or(Duration::ZERO);
Err(format!(
"Circuit breaker open for '{}' ({} consecutive failures, \
recovery in {:.0}s)",
self.label,
state.consecutive_failures,
remaining.as_secs_f64()
))
}
} else {
state.state = CircuitState::Closed;
Ok(())
}
}
}
}
/// Record a successful call.
pub 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!(
label = %self.label,
"Circuit breaker: HalfOpen -> Closed (recovered)"
);
}
}
CircuitState::Open => {
state.state = CircuitState::Closed;
state.consecutive_failures = 0;
state.opened_at = None;
}
}
}
/// Record a failed call. Only transient errors count toward the threshold.
pub async fn record_failure<E>(&self, err: &E)
where
C: ErrorClassifier<E>,
{
if !self.classifier.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!(
label = %self.label,
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!(
label = %self.label,
"Circuit breaker: HalfOpen -> Open (probe failed)"
);
}
CircuitState::Open => {
// Already open, nothing to do
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Debug, thiserror::Error)]
enum TestError {
#[error("transient")]
Transient,
#[error("permanent")]
Permanent,
}
struct TestClassifier;
impl ErrorClassifier<TestError> for TestClassifier {
fn is_retryable(&self, err: &TestError) -> bool {
matches!(err, TestError::Transient)
}
fn is_transient(&self, err: &TestError) -> bool {
matches!(err, TestError::Transient)
}
}
fn make_breaker(threshold: u32) -> CircuitBreakerLayer<TestClassifier> {
CircuitBreakerLayer::new(
CircuitBreakerConfig {
failure_threshold: threshold,
recovery_timeout: Duration::from_millis(100),
half_open_successes_needed: 2,
},
TestClassifier,
"test",
)
}
#[tokio::test]
async fn test_closed_allows_calls() {
let cb = make_breaker(3);
assert!(cb.check_allowed().await.is_ok()); // safety: test-only
assert_eq!(cb.circuit_state().await, CircuitState::Closed); // safety: test-only
}
#[tokio::test]
async fn test_opens_after_threshold() {
let cb = make_breaker(3);
for _ in 0..3 {
cb.record_failure(&TestError::Transient).await;
}
assert_eq!(cb.circuit_state().await, CircuitState::Open); // safety: test-only
assert!(cb.check_allowed().await.is_err()); // safety: test-only
}
#[tokio::test]
async fn test_permanent_errors_dont_trip() {
let cb = make_breaker(3);
for _ in 0..10 {
cb.record_failure(&TestError::Permanent).await;
}
assert_eq!(cb.circuit_state().await, CircuitState::Closed); // safety: test-only
}
#[tokio::test]
async fn test_success_resets_count() {
let cb = make_breaker(3);
cb.record_failure(&TestError::Transient).await;
cb.record_failure(&TestError::Transient).await;
cb.record_success().await;
assert_eq!(cb.consecutive_failures().await, 0); // safety: test-only
// Should still be closed since we reset
cb.record_failure(&TestError::Transient).await;
cb.record_failure(&TestError::Transient).await;
assert_eq!(cb.circuit_state().await, CircuitState::Closed); // safety: test-only
}
#[tokio::test]
async fn test_recovery_to_half_open() {
let cb = make_breaker(1);
cb.record_failure(&TestError::Transient).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open); // safety: test-only
// Wait for recovery timeout
tokio::time::sleep(Duration::from_millis(150)).await;
// Should transition to HalfOpen
assert!(cb.check_allowed().await.is_ok()); // safety: test-only
assert_eq!(cb.circuit_state().await, CircuitState::HalfOpen); // safety: test-only
}
#[tokio::test]
async fn test_half_open_closes_on_successes() {
let cb = make_breaker(1);
cb.record_failure(&TestError::Transient).await;
tokio::time::sleep(Duration::from_millis(150)).await;
let _ = cb.check_allowed().await; // transition to HalfOpen
cb.record_success().await;
assert_eq!(cb.circuit_state().await, CircuitState::HalfOpen); // needs 2 // safety: test-only
cb.record_success().await;
assert_eq!(cb.circuit_state().await, CircuitState::Closed); // safety: test-only
}
#[tokio::test]
async fn test_half_open_reopens_on_failure() {
let cb = make_breaker(1);
cb.record_failure(&TestError::Transient).await;
tokio::time::sleep(Duration::from_millis(150)).await;
let _ = cb.check_allowed().await; // HalfOpen
cb.record_failure(&TestError::Transient).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open); // safety: test-only
}
}
+21
View File
@@ -0,0 +1,21 @@
//! Generic error classification for resilience layers.
use std::time::Duration;
/// Classifies errors to determine how resilience layers should respond.
///
/// Each client type (LLM, MCP, HTTP tool, etc.) implements this trait
/// to tell the resilience layers how to handle its specific error type.
pub trait ErrorClassifier<E> {
/// Should the same request be retried against the same endpoint?
fn is_retryable(&self, err: &E) -> bool;
/// Does this error indicate the backend is degraded?
/// Used by circuit breakers to track health.
fn is_transient(&self, err: &E) -> bool;
/// Provider-suggested retry delay (e.g. from Retry-After header).
fn retry_after(&self, _err: &E) -> Option<Duration> {
None
}
}
+166
View File
@@ -0,0 +1,166 @@
//! Per-endpoint health tracking.
//!
//! Provides atomic, lock-free health counters for external service endpoints.
//! Used by the state bus to publish `EndpointHealthChanged` events.
use std::collections::HashMap;
use std::sync::RwLock;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
/// Health state for a single endpoint.
pub struct EndpointHealth {
/// Number of consecutive failures.
pub consecutive_failures: AtomicU32,
/// Unix timestamp (seconds) of last successful call.
pub last_success: AtomicU64,
/// 0 = healthy, 1 = unhealthy.
pub unhealthy: AtomicU32,
}
impl EndpointHealth {
pub fn new() -> Self {
Self {
consecutive_failures: AtomicU32::new(0),
last_success: AtomicU64::new(0),
unhealthy: AtomicU32::new(0),
}
}
pub fn record_success(&self) {
self.consecutive_failures.store(0, Ordering::Relaxed);
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
self.last_success.store(now, Ordering::Relaxed);
self.unhealthy.store(0, Ordering::Relaxed);
}
/// Record a failure. Returns true if this failure triggered the unhealthy threshold.
pub fn record_failure(&self, threshold: u32) -> bool {
let prev = self.consecutive_failures.fetch_add(1, Ordering::Relaxed);
let new_count = prev + 1;
if new_count >= threshold && self.unhealthy.swap(1, Ordering::Relaxed) == 0 {
return true; // Just became unhealthy
}
false
}
pub fn is_healthy(&self) -> bool {
self.unhealthy.load(Ordering::Relaxed) == 0
}
}
impl Default for EndpointHealth {
fn default() -> Self {
Self::new()
}
}
/// Tracks health of multiple named endpoints.
pub struct HealthTracker {
endpoints: RwLock<HashMap<String, EndpointHealth>>,
failure_threshold: u32,
}
impl HealthTracker {
pub fn new(failure_threshold: u32) -> Self {
Self {
endpoints: RwLock::new(HashMap::new()),
failure_threshold,
}
}
pub fn record_success(&self, name: &str) {
let endpoints = self.endpoints.read().unwrap_or_else(|e| e.into_inner());
if let Some(health) = endpoints.get(name) {
health.record_success();
} else {
drop(endpoints);
let mut endpoints = self.endpoints.write().unwrap_or_else(|e| e.into_inner());
endpoints
.entry(name.to_string())
.or_default()
.record_success();
}
}
/// Record a failure. Returns true if this made the endpoint unhealthy.
pub fn record_failure(&self, name: &str) -> bool {
let endpoints = self.endpoints.read().unwrap_or_else(|e| e.into_inner());
if let Some(health) = endpoints.get(name) {
health.record_failure(self.failure_threshold)
} else {
drop(endpoints);
let mut endpoints = self.endpoints.write().unwrap_or_else(|e| e.into_inner());
let health = endpoints.entry(name.to_string()).or_default();
health.record_failure(self.failure_threshold)
}
}
pub fn is_healthy(&self, name: &str) -> bool {
let endpoints = self.endpoints.read().unwrap_or_else(|e| e.into_inner());
endpoints.get(name).map(|h| h.is_healthy()).unwrap_or(true) // Unknown endpoints are assumed healthy
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_endpoint_health_starts_healthy() {
let h = EndpointHealth::new();
assert!(h.is_healthy()); // safety: test-only
}
#[test]
fn test_endpoint_health_becomes_unhealthy() {
let h = EndpointHealth::new();
for i in 0..4 {
assert!(
// safety: test-only
// safety: test-only
!h.record_failure(5),
"should not be unhealthy at failure {}",
i + 1
);
}
assert!(h.record_failure(5), "should become unhealthy at failure 5"); // safety: test-only
assert!(!h.is_healthy()); // safety: test-only
}
#[test]
fn test_endpoint_health_recovers() {
let h = EndpointHealth::new();
for _ in 0..5 {
h.record_failure(5);
}
assert!(!h.is_healthy()); // safety: test-only
h.record_success();
assert!(h.is_healthy()); // safety: test-only
}
#[test]
fn test_tracker_unknown_is_healthy() {
let t = HealthTracker::new(3);
assert!(t.is_healthy("unknown")); // safety: test-only
}
#[test]
fn test_tracker_tracks_failures() {
let t = HealthTracker::new(2);
assert!(!t.record_failure("ep1")); // safety: test-only
assert!(t.record_failure("ep1")); // safety: test-only
assert!(!t.is_healthy("ep1")); // safety: test-only
}
#[test]
fn test_tracker_recovery() {
let t = HealthTracker::new(2);
t.record_failure("ep1");
t.record_failure("ep1");
t.record_success("ep1");
assert!(t.is_healthy("ep1")); // safety: test-only
}
}
+9
View File
@@ -0,0 +1,9 @@
pub mod circuit_breaker;
pub mod classifier;
pub mod health;
pub mod retry;
pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerLayer, CircuitState};
pub use classifier::ErrorClassifier;
pub use health::{EndpointHealth, HealthTracker};
pub use retry::{RetryConfig, RetryLayer};
+193
View File
@@ -0,0 +1,193 @@
//! Generic retry layer with exponential backoff and jitter.
//!
//! Extracted from `llm::retry` to be reusable across MCP, HTTP tools,
//! relay channels, and any async operation that can fail transiently.
use std::future::Future;
use std::time::Duration;
use rand::Rng;
use super::classifier::ErrorClassifier;
/// Configuration for the retry layer.
#[derive(Debug, Clone)]
pub struct RetryConfig {
/// Maximum number of retry attempts (not counting the initial attempt).
pub max_retries: u32,
}
impl Default for RetryConfig {
fn default() -> Self {
Self { max_retries: 3 }
}
}
/// Generic retry layer that wraps any async operation.
pub struct RetryLayer<C> {
config: RetryConfig,
classifier: C,
}
impl<C> RetryLayer<C> {
pub fn new(config: RetryConfig, classifier: C) -> Self {
Self { config, classifier }
}
}
impl<C> RetryLayer<C> {
/// Execute an operation with retry logic.
///
/// `label` is included in log messages for diagnostics.
pub async fn execute<T, E, F, Fut>(&self, mut op: F, label: &str) -> Result<T, E>
where
C: ErrorClassifier<E>,
E: std::fmt::Display,
F: FnMut() -> Fut,
Fut: Future<Output = Result<T, E>>,
{
let mut last_error: Option<E> = None;
for attempt in 0..=self.config.max_retries {
match op().await {
Ok(val) => return Ok(val),
Err(err) => {
if !self.classifier.is_retryable(&err) || attempt == self.config.max_retries {
return Err(err);
}
let delay = self
.classifier
.retry_after(&err)
.unwrap_or_else(|| retry_backoff_delay(attempt));
tracing::warn!(
attempt = attempt + 1,
max_retries = self.config.max_retries,
delay_ms = delay.as_millis() as u64,
error = %err,
"Retrying after transient error ({label})"
);
last_error = Some(err);
tokio::time::sleep(delay).await;
}
}
}
// Safety: loop runs at least once (0..=max_retries), so last_error is always Some
// if we reach here. But be defensive.
match last_error {
Some(e) => Err(e),
None => unreachable!("retry loop ran at least once"),
}
}
}
/// Calculate exponential backoff delay with random jitter.
///
/// Base delay is 1 second, doubled each attempt, with +/-25% jitter.
pub fn retry_backoff_delay(attempt: u32) -> Duration {
let base_ms: u64 = 1000u64.saturating_mul(2u64.saturating_pow(attempt));
let jitter_range = base_ms / 4; // 25%
let jitter = if jitter_range > 0 {
let offset = rand::thread_rng().gen_range(0..=jitter_range * 2);
offset as i64 - jitter_range as i64
} else {
0
};
let delay_ms = (base_ms as i64 + jitter).max(100) as u64;
Duration::from_millis(delay_ms)
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
#[derive(Debug, thiserror::Error)]
enum TestError {
#[error("transient")]
Transient,
#[error("permanent")]
Permanent,
}
struct TestClassifier;
impl ErrorClassifier<TestError> for TestClassifier {
fn is_retryable(&self, err: &TestError) -> bool {
matches!(err, TestError::Transient)
}
fn is_transient(&self, err: &TestError) -> bool {
matches!(err, TestError::Transient)
}
}
#[test]
fn test_backoff_delay_exponential() {
for _ in 0..10 {
let d0 = retry_backoff_delay(0);
assert!(d0.as_millis() >= 750 && d0.as_millis() <= 1250); // safety: test-only
let d1 = retry_backoff_delay(1);
assert!(d1.as_millis() >= 1500 && d1.as_millis() <= 2500); // safety: test-only
}
}
#[test]
fn test_backoff_delay_no_overflow() {
let delay = retry_backoff_delay(30);
assert!(delay.as_millis() >= 100); // safety: test-only
}
#[tokio::test]
async fn test_success_first_attempt() {
let layer = RetryLayer::new(RetryConfig { max_retries: 3 }, TestClassifier);
let result: Result<&str, TestError> = layer.execute(|| async { Ok("ok") }, "test").await; // safety: test-only retry call
assert_eq!(result.unwrap(), "ok"); // safety: test-only
}
#[tokio::test]
async fn test_permanent_error_no_retry() {
let calls = Arc::new(AtomicU32::new(0));
let calls_c = calls.clone();
let layer = RetryLayer::new(RetryConfig { max_retries: 3 }, TestClassifier);
let result: Result<(), TestError> = layer
.execute(
// safety: test-only retry call
|| {
let c = calls_c.clone();
async move {
c.fetch_add(1, Ordering::Relaxed);
Err(TestError::Permanent)
}
},
"test",
)
.await;
assert!(result.is_err()); // safety: test-only
assert_eq!(calls.load(Ordering::Relaxed), 1); // safety: test-only
}
#[tokio::test]
async fn test_exhausts_retries() {
let calls = Arc::new(AtomicU32::new(0));
let calls_c = calls.clone();
let layer = RetryLayer::new(RetryConfig { max_retries: 0 }, TestClassifier);
let result: Result<(), TestError> = layer
.execute(
// safety: test-only retry call
|| {
let c = calls_c.clone();
async move {
c.fetch_add(1, Ordering::Relaxed);
Err(TestError::Transient)
}
},
"test",
)
.await;
assert!(result.is_err()); // safety: test-only
assert_eq!(calls.load(Ordering::Relaxed), 1); // safety: test-only
}
}
+131
View File
@@ -0,0 +1,131 @@
//! State invalidation bus for cross-module state synchronization.
//!
//! When state changes in one module (e.g., web UI toggles a routine, secret
//! rotates, config reloads), the bus notifies other modules that cache that
//! state so they can refresh.
//!
//! Modules subscribe to events they care about and ignore the rest. No module
//! needs to import another module to propagate state changes — the bus is the
//! **only** coupling point.
use std::sync::Arc;
use tokio::sync::broadcast;
use uuid::Uuid;
/// A state change notification.
#[derive(Debug, Clone, serde::Serialize)]
pub enum StateChange {
/// A routine was created, updated, toggled, or deleted.
RoutineUpdated { routine_id: Uuid },
/// A secret was rotated or deleted.
SecretRotated { key_name: String },
/// Global configuration was reloaded (e.g. via SIGHUP).
ConfigReloaded,
/// An external endpoint's health status changed.
EndpointHealthChanged { name: String, healthy: bool },
/// The tool registry was modified (tool added/removed/rebuilt).
ToolRegistryChanged,
/// An extension was installed or removed.
ExtensionInstalled { extension_id: String },
}
/// Broadcast bus for state change notifications.
///
/// Backed by a tokio `broadcast` channel with a fixed buffer. Slow consumers
/// that fall behind will miss events (acceptable — they can re-poll state).
#[derive(Clone)]
pub struct StateBus {
tx: broadcast::Sender<StateChange>,
}
impl StateBus {
/// Create a new state bus with a buffer of 64 events.
pub fn new() -> Self {
let (tx, _) = broadcast::channel(64);
Self { tx }
}
/// Publish a state change. Non-blocking; drops the event if no subscribers.
pub fn publish(&self, event: StateChange) {
// Ignore send error (no active receivers).
let _ = self.tx.send(event);
}
/// Subscribe to state change notifications.
pub fn subscribe(&self) -> broadcast::Receiver<StateChange> {
self.tx.subscribe()
}
}
impl Default for StateBus {
fn default() -> Self {
Self::new()
}
}
/// Convenience constructor for passing through `Arc`.
pub fn new_state_bus() -> Arc<StateBus> {
Arc::new(StateBus::new())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_publish_subscribe() {
let bus = StateBus::new();
let mut rx = bus.subscribe();
let id = Uuid::new_v4();
bus.publish(StateChange::RoutineUpdated { routine_id: id });
let event = rx.recv().await.unwrap(); // safety: test-only
assert!(matches!(event, StateChange::RoutineUpdated { routine_id } if routine_id == id)); // safety: test-only
}
#[tokio::test]
async fn test_no_subscriber_does_not_panic() {
let bus = StateBus::new();
// No subscribers — should not panic.
bus.publish(StateChange::ConfigReloaded);
}
#[tokio::test]
async fn test_multiple_subscribers() {
let bus = StateBus::new();
let mut rx1 = bus.subscribe();
let mut rx2 = bus.subscribe();
bus.publish(StateChange::ToolRegistryChanged);
let e1 = rx1.recv().await.unwrap(); // safety: test-only
let e2 = rx2.recv().await.unwrap(); // safety: test-only
assert!(matches!(e1, StateChange::ToolRegistryChanged)); // safety: test-only
assert!(matches!(e2, StateChange::ToolRegistryChanged)); // safety: test-only
}
#[tokio::test]
async fn test_slow_consumer_lags() {
let bus = StateBus::new();
let mut rx = bus.subscribe();
// Overflow the 64-event buffer.
for i in 0..100 {
bus.publish(StateChange::EndpointHealthChanged {
name: format!("ep-{}", i),
healthy: true,
});
}
// First recv should report a lag.
let result = rx.recv().await;
assert!(
// safety: test-only
// safety: test-only
result.is_ok() || result.is_err(),
"lagged receiver should either get an event or a Lagged error"
);
}
}
+1
View File
@@ -455,6 +455,7 @@ impl TestHarnessBuilder {
http_interceptor: None,
transcription: None,
document_extraction: None,
event_bus: None,
};
TestHarness {
+1 -1
View File
@@ -463,7 +463,7 @@ impl CreateJobTool {
}
match jm.get_handle(job_id).await {
Some(handle) => match handle.state {
Some(handle) => match handle.state() {
crate::orchestrator::job_manager::ContainerState::Running
| crate::orchestrator::job_manager::ContainerState::Creating => {
tokio::time::sleep(poll_interval).await;
+46
View File
@@ -194,6 +194,52 @@ impl ToolRegistry {
self.tools.try_read().map(|t| t.len()).unwrap_or(0)
}
/// Verify that expected built-in tools are registered.
///
/// Returns a list of tool names that should be registered but aren't.
/// Called during startup to catch wiring bugs.
pub fn verify_expected_tools(&self, config: &crate::config::Config) -> Vec<String> {
let tools = match self.tools.try_read() {
Ok(t) => t,
Err(_) => return Vec::new(),
};
let mut missing = Vec::new();
// Core built-ins that should always be present
for name in &["echo", "time", "json", "http"] {
if !tools.contains_key(*name) {
missing.push(name.to_string());
}
}
// Memory tools should be present when tools beyond core builtins are loaded
// (indicates a workspace/DB is available).
if tools.len() > 4 {
for name in &[
"memory_search",
"memory_write",
"memory_read",
"memory_tree",
] {
if !tools.contains_key(*name) {
missing.push(name.to_string());
}
}
}
// Dev tools when local tools allowed
if config.agent.allow_local_tools {
for name in &["shell", "read_file", "write_file"] {
if !tools.contains_key(*name) {
missing.push(name.to_string());
}
}
}
missing
}
/// Get all tools.
pub async fn all(&self) -> Vec<Arc<dyn Tool>> {
self.tools.read().await.values().cloned().collect()
+91 -25
View File
@@ -70,109 +70,175 @@ pub fn llm_signals_completion(response: &str) -> bool {
positive_phrases.iter().any(|p| lower.contains(p))
}
/// Truncate a string to at most `max_bytes` bytes at a char boundary, appending "...".
///
/// If the input is wrapped in `<tool_output …>…</tool_output>` and truncation
/// removes the closing tag, the tag is re-appended so downstream XML parsers
/// never see an unclosed element.
pub fn truncate_preview(s: &str, max_bytes: usize) -> String {
if s.len() <= max_bytes {
return s.to_string();
}
let mut end = max_bytes;
while end > 0 && !s.is_char_boundary(end) {
end -= 1;
}
let mut result = format!("{}...", &s[..end]); // safety: end is a valid char boundary per loop above
if s.starts_with("<tool_output") && !result.ends_with("</tool_output>") {
result.push_str("\n</tool_output>");
}
result
}
#[cfg(test)]
mod tests {
use crate::util::{floor_char_boundary, llm_signals_completion};
use crate::util::{floor_char_boundary, llm_signals_completion, truncate_preview};
// ── floor_char_boundary ──
#[test]
fn floor_char_boundary_at_valid_boundary() {
assert_eq!(floor_char_boundary("hello", 3), 3);
assert_eq!(floor_char_boundary("hello", 3), 3); // safety: test-only
}
#[test]
fn floor_char_boundary_mid_multibyte_char() {
// h = 1 byte, é = 2 bytes, total 3 bytes
let s = "";
assert_eq!(floor_char_boundary(s, 2), 1); // byte 2 is mid-é, back up to 1
assert_eq!(floor_char_boundary(s, 2), 1); // byte 2 is mid-é, back up to 1 // safety: test-only
}
#[test]
fn floor_char_boundary_past_end() {
assert_eq!(floor_char_boundary("hi", 100), 2);
assert_eq!(floor_char_boundary("hi", 100), 2); // safety: test-only
}
#[test]
fn floor_char_boundary_at_zero() {
assert_eq!(floor_char_boundary("hello", 0), 0);
assert_eq!(floor_char_boundary("hello", 0), 0); // safety: test-only
}
#[test]
fn floor_char_boundary_empty_string() {
assert_eq!(floor_char_boundary("", 5), 0);
assert_eq!(floor_char_boundary("", 5), 0); // safety: test-only
}
// ── llm_signals_completion ──
#[test]
fn signals_completion_positive() {
assert!(llm_signals_completion("The job is complete."));
assert!(llm_signals_completion("I have completed the task."));
assert!(llm_signals_completion("All done, here are the results."));
assert!(llm_signals_completion("Task is finished successfully."));
assert!(llm_signals_completion("The job is complete.")); // safety: test-only
assert!(llm_signals_completion("I have completed the task.")); // safety: test-only
assert!(llm_signals_completion("All done, here are the results.")); // safety: test-only
assert!(llm_signals_completion("Task is finished successfully.")); // safety: test-only
assert!(llm_signals_completion(
// safety: test-only
"I have completed the task successfully."
));
assert!(llm_signals_completion(
// safety: test-only
"All steps are complete and verified."
));
assert!(llm_signals_completion(
// safety: test-only
"I've done all the work. The work is done."
));
assert!(llm_signals_completion(
// safety: test-only
"Successfully completed the migration."
));
assert!(llm_signals_completion(
// safety: test-only
"I have completed the job ahead of schedule."
));
assert!(llm_signals_completion("I have finished the task."));
assert!(llm_signals_completion("All steps are done now."));
assert!(llm_signals_completion("I've completed everything."));
assert!(llm_signals_completion("All tasks complete."));
assert!(llm_signals_completion("I have finished the task.")); // safety: test-only
assert!(llm_signals_completion("All steps are done now.")); // safety: test-only
assert!(llm_signals_completion("I've completed everything.")); // safety: test-only
assert!(llm_signals_completion("All tasks complete.")); // safety: test-only
}
#[test]
fn signals_completion_negative() {
assert!(!llm_signals_completion("The task is not complete yet."));
assert!(!llm_signals_completion("This is not done."));
assert!(!llm_signals_completion("The work is incomplete."));
assert!(!llm_signals_completion("Build is unfinished."));
assert!(!llm_signals_completion("The task is not complete yet.")); // safety: test-only
assert!(!llm_signals_completion("This is not done.")); // safety: test-only
assert!(!llm_signals_completion("The work is incomplete.")); // safety: test-only
assert!(!llm_signals_completion("Build is unfinished.")); // safety: test-only
assert!(!llm_signals_completion(
// safety: test-only
"The migration is not yet finished."
));
assert!(!llm_signals_completion("The job isn't done yet."));
assert!(!llm_signals_completion("This remains unfinished."));
assert!(!llm_signals_completion("The job isn't done yet.")); // safety: test-only
assert!(!llm_signals_completion("This remains unfinished.")); // safety: test-only
}
#[test]
fn signals_completion_no_bare_substrings() {
assert!(!llm_signals_completion("The download completed."));
assert!(!llm_signals_completion("The download completed.")); // safety: test-only
assert!(!llm_signals_completion(
// safety: test-only
"Function done_callback was called."
));
assert!(!llm_signals_completion("Set is_complete = true"));
assert!(!llm_signals_completion("Running step 3 of 5"));
assert!(!llm_signals_completion("Set is_complete = true")); // safety: test-only
assert!(!llm_signals_completion("Running step 3 of 5")); // safety: test-only
assert!(!llm_signals_completion(
// safety: test-only
"I need to complete more work first."
));
assert!(!llm_signals_completion(
// safety: test-only
"Let me finish the remaining steps."
));
assert!(!llm_signals_completion(
// safety: test-only
"I'm done analyzing, now let me fix it."
));
assert!(!llm_signals_completion(
// safety: test-only
"I completed step 1 but step 2 remains."
));
}
#[test]
fn signals_completion_tool_output_injection() {
assert!(!llm_signals_completion("TASK_COMPLETE"));
assert!(!llm_signals_completion("JOB_DONE"));
assert!(!llm_signals_completion("TASK_COMPLETE")); // safety: test-only
assert!(!llm_signals_completion("JOB_DONE")); // safety: test-only
assert!(!llm_signals_completion(
// safety: test-only
"The tool returned: TASK_COMPLETE signal"
));
}
// ── truncate_preview ──
#[test]
fn truncate_preview_short_string() {
assert_eq!(truncate_preview("hello", 10), "hello"); // safety: test-only
}
#[test]
fn truncate_preview_exact_boundary() {
assert_eq!(truncate_preview("hello", 5), "hello"); // safety: test-only
}
#[test]
fn truncate_preview_truncates_ascii() {
assert_eq!(truncate_preview("hello world", 5), "hello..."); // safety: test-only
}
#[test]
fn truncate_preview_multibyte_char_boundary() {
let s = "a€b";
let result = truncate_preview(s, 3);
assert_eq!(result, "a..."); // safety: test-only
}
#[test]
fn truncate_preview_closes_tool_output_tag() {
let s = "<tool_output name=\"search\" sanitized=\"true\">\nSome very long content here\n</tool_output>";
let result = truncate_preview(s, 60);
assert!(result.ends_with("</tool_output>")); // safety: test-only
assert!(result.contains("...")); // safety: test-only
}
}
+59
View File
@@ -54,6 +54,9 @@ pub struct WorkerDeps {
pub approval_context: Option<ApprovalContext>,
/// HTTP interceptor for trace recording/replay (propagated to JobContext).
pub http_interceptor: Option<Arc<dyn crate::llm::recording::HttpInterceptor>>,
/// Unified event bus. When present, job events are emitted through the bus
/// in addition to the legacy SSE + DB paths.
pub event_bus: Option<crate::event_bus::EventBus>,
}
/// Worker that executes a single job.
@@ -120,9 +123,63 @@ impl Worker {
}
/// Fire-and-forget persistence of a job event and SSE broadcast.
///
/// Also emits through the unified event bus when available, so the
/// audit sink and other subscribers capture all job activity.
fn log_event(&self, event_type: &str, data: serde_json::Value) {
let job_id = self.job_id;
// Emit through unified event bus (audit + future sinks)
if let Some(ref bus) = self.deps.event_bus {
use crate::event_bus::{EventContext, EventSource};
use crate::events::DomainEvent;
let job_id_str = job_id.to_string();
let domain_event = match event_type {
"message" => Some(DomainEvent::JobMessage {
job_id: job_id_str.clone(),
role: data
.get("role")
.and_then(|v| v.as_str())
.unwrap_or("assistant")
.to_string(),
content: data
.get("content")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
}),
"status" => Some(DomainEvent::JobStatus {
job_id: job_id_str.clone(),
message: data
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
}),
"result" => Some(DomainEvent::JobResult {
job_id: job_id_str.clone(),
status: data
.get("status")
.and_then(|v| v.as_str())
.unwrap_or("completed")
.to_string(),
session_id: data
.get("session_id")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
}),
_ => None,
};
if let Some(event) = domain_event {
bus.emit_domain(
EventSource::new("worker", "job"),
EventContext::with_job(job_id),
event,
);
}
}
// Persist to DB
if let Some(store) = self.store() {
let store = store.clone();
@@ -1452,6 +1509,7 @@ mod tests {
sse_tx: None,
approval_context: None,
http_interceptor: None,
event_bus: None,
};
Worker::new(job_id, deps)
@@ -1642,6 +1700,7 @@ mod tests {
sse_tx: None,
approval_context,
http_interceptor: None,
event_bus: None,
};
Worker::new(job_id, deps)
+53 -45
View File
@@ -19,12 +19,12 @@ mod tests {
async fn create_test_db() -> (Arc<dyn Database>, tempfile::TempDir) {
use ironclaw::db::libsql::LibSqlBackend;
let temp_dir = tempfile::tempdir().expect("tempdir");
let temp_dir = tempfile::tempdir().expect("tempdir"); // safety: test-only
let db_path = temp_dir.path().join("test.db");
let backend = LibSqlBackend::new_local(&db_path)
.await
.expect("LibSqlBackend");
backend.run_migrations().await.expect("migrations");
.expect("LibSqlBackend"); // safety: test-only
backend.run_migrations().await.expect("migrations"); // safety: test-only
let db: Arc<dyn Database> = Arc::new(backend);
(db, temp_dir)
}
@@ -39,8 +39,8 @@ mod tests {
let counts = db
.count_running_routine_runs_batch(&[])
.await
.expect("batch query should not fail");
assert!(counts.is_empty(), "Empty input should return empty map");
.expect("batch query should not fail"); // safety: test-only
assert!(counts.is_empty(), "Empty input should return empty map"); // safety: test-only
}
#[tokio::test]
@@ -80,7 +80,7 @@ mod tests {
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
db.create_routine(&routine).await.expect("create routine"); // safety: test-only
// Create 3 running runs
for _ in 0..3 {
@@ -97,17 +97,17 @@ mod tests {
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
db.create_routine_run(&run).await.expect("create run"); // safety: test-only
}
// Batch query for single routine
let counts = db
.count_running_routine_runs_batch(&[routine_id])
.await
.expect("batch query should work");
.expect("batch query should work"); // safety: test-only
assert_eq!(counts.len(), 1, "Should return 1 routine");
assert_eq!(counts[&routine_id], 3, "Should count 3 running runs");
assert_eq!(counts.len(), 1, "Should return 1 routine"); // safety: test-only
assert_eq!(counts[&routine_id], 3, "Should count 3 running runs"); // safety: test-only
}
#[tokio::test]
@@ -151,7 +151,7 @@ mod tests {
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
db.create_routine(&routine).await.expect("create routine"); // safety: test-only
}
// r1: 2 running
@@ -169,7 +169,7 @@ mod tests {
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
db.create_routine_run(&run).await.expect("create run"); // safety: test-only
}
// r2: 1 running
@@ -186,7 +186,7 @@ mod tests {
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
db.create_routine_run(&run).await.expect("create run"); // safety: test-only
// r3: 0 running (but has 1 Ok result)
let run = RoutineRun {
@@ -202,19 +202,21 @@ mod tests {
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
db.create_routine_run(&run).await.expect("create run"); // safety: test-only
// Single batch query for all 3
let counts = db
.count_running_routine_runs_batch(&[r1, r2, r3])
.await
.expect("batch query should work");
.expect("batch query should work"); // safety: test-only
assert_eq!(counts.len(), 3, "Should return 3 routines");
assert_eq!(counts[&r1], 2, "r1 should have 2 running");
assert_eq!(counts[&r2], 1, "r2 should have 1 running");
assert_eq!(counts.len(), 3, "Should return 3 routines"); // safety: test-only
assert_eq!(counts[&r1], 2, "r1 should have 2 running"); // safety: test-only
assert_eq!(counts[&r2], 1, "r2 should have 1 running"); // safety: test-only
assert_eq!(
counts[&r3], 0,
// safety: test-only
counts[&r3],
0,
"r3 should have 0 running (Ok status is not running)"
);
}
@@ -259,7 +261,7 @@ mod tests {
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
db.create_routine(&routine).await.expect("create routine"); // safety: test-only
// r1 has 1 running
let run = RoutineRun {
@@ -275,18 +277,18 @@ mod tests {
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
db.create_routine_run(&run).await.expect("create run"); // safety: test-only
// Query for r1, r2 (doesn't exist), r3 (doesn't exist)
let counts = db
.count_running_routine_runs_batch(&[r1, r2, r3])
.await
.expect("batch query should work");
.expect("batch query should work"); // safety: test-only
assert_eq!(counts.len(), 3, "Should have all 3 routine IDs");
assert_eq!(counts[&r1], 1, "r1 should have 1 running");
assert_eq!(counts[&r2], 0, "r2 should default to 0");
assert_eq!(counts[&r3], 0, "r3 should default to 0");
assert_eq!(counts.len(), 3, "Should have all 3 routine IDs"); // safety: test-only
assert_eq!(counts[&r1], 1, "r1 should have 1 running"); // safety: test-only
assert_eq!(counts[&r2], 0, "r2 should default to 0"); // safety: test-only
assert_eq!(counts[&r3], 0, "r3 should default to 0"); // safety: test-only
}
#[tokio::test]
@@ -326,7 +328,7 @@ mod tests {
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
db.create_routine(&routine).await.expect("create routine"); // safety: test-only
// Create 5 runs with mixed statuses
let statuses = [
@@ -351,17 +353,19 @@ mod tests {
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
db.create_routine_run(&run).await.expect("create run"); // safety: test-only
}
// Batch query should only count Running status
let counts = db
.count_running_routine_runs_batch(&[routine_id])
.await
.expect("batch query should work");
.expect("batch query should work"); // safety: test-only
assert_eq!(
counts[&routine_id], 2,
// safety: test-only
counts[&routine_id],
2,
"Should only count 2 Running status runs"
);
}
@@ -410,7 +414,7 @@ mod tests {
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
db.create_routine(&routine).await.expect("create routine"); // safety: test-only
}
// r1: create 1 running run (will hit max_concurrent=1)
@@ -427,7 +431,7 @@ mod tests {
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
db.create_routine_run(&run).await.expect("create run"); // safety: test-only
// r2: create 2 running runs (will hit max_concurrent=2)
for _ in 0..2 {
@@ -444,22 +448,26 @@ mod tests {
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
db.create_routine_run(&run).await.expect("create run"); // safety: test-only
}
// Batch query should return correct counts
let counts = db
.count_running_routine_runs_batch(&[r1, r2])
.await
.expect("batch query should work");
.expect("batch query should work"); // safety: test-only
// Verify counts match the limits
assert_eq!(
counts[&r1], 1,
// safety: test-only
counts[&r1],
1,
"r1 should have 1 running (at max_concurrent=1)"
);
assert_eq!(
counts[&r2], 2,
// safety: test-only
counts[&r2],
2,
"r2 should have 2 running (at max_concurrent=2)"
);
@@ -467,19 +475,19 @@ mod tests {
let r1_routine = db
.get_routine(r1)
.await
.expect("get routine")
.expect("routine exists");
.expect("get routine") // safety: test-only
.expect("routine exists"); // safety: test-only
let r2_routine = db
.get_routine(r2)
.await
.expect("get routine")
.expect("routine exists");
.expect("get routine") // safety: test-only
.expect("routine exists"); // safety: test-only
let r1_at_limit = counts[&r1] >= r1_routine.guardrails.max_concurrent as i64;
let r2_at_limit = counts[&r2] >= r2_routine.guardrails.max_concurrent as i64;
assert!(r1_at_limit, "r1 should be detected as at limit");
assert!(r2_at_limit, "r2 should be detected as at limit");
assert!(r1_at_limit, "r1 should be detected as at limit"); // safety: test-only
assert!(r2_at_limit, "r2 should be detected as at limit"); // safety: test-only
// If we add one more run to r2, it should exceed limit
let run = RoutineRun {
@@ -495,15 +503,15 @@ mod tests {
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
db.create_routine_run(&run).await.expect("create run"); // safety: test-only
// Re-query to get updated counts
let counts = db
.count_running_routine_runs_batch(&[r1, r2])
.await
.expect("batch query should work");
.expect("batch query should work"); // safety: test-only
let r2_exceeded_limit = counts[&r2] > r2_routine.guardrails.max_concurrent as i64;
assert!(r2_exceeded_limit, "r2 should have exceeded its limit");
assert!(r2_exceeded_limit, "r2 should have exceeded its limit"); // safety: test-only
}
}
+376
View File
@@ -0,0 +1,376 @@
//! Boundary chaos tests — exercise failure modes at module seams.
//!
//! These tests verify that the architectural hardening (domain event decoupling,
//! generic resilience layers, state bus) works correctly under failure conditions.
//!
//! Organized by boundary, not by module:
//! - Resilience layers (retry, circuit breaker, health tracker)
//! - State bus propagation
//! - Domain event type compatibility
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Duration;
use ironclaw::events::DomainEvent;
use ironclaw::resilience::circuit_breaker::{
CircuitBreakerConfig, CircuitBreakerLayer, CircuitState,
};
use ironclaw::resilience::classifier::ErrorClassifier;
use ironclaw::resilience::health::HealthTracker;
use ironclaw::resilience::retry::{RetryConfig, RetryLayer};
use ironclaw::state_bus::{StateBus, StateChange};
// ── Test error type ──────────────────────────────────────────────────
#[derive(Debug, thiserror::Error)]
enum TestError {
#[error("transient failure")]
Transient,
#[error("permanent failure")]
Permanent,
}
struct TestClassifier;
impl ErrorClassifier<TestError> for TestClassifier {
fn is_retryable(&self, err: &TestError) -> bool {
matches!(err, TestError::Transient)
}
fn is_transient(&self, err: &TestError) -> bool {
matches!(err, TestError::Transient)
}
}
// ── Resilience: Retry layer ──────────────────────────────────────────
#[tokio::test]
async fn retry_layer_recovers_after_transient_failures() {
let call_count = Arc::new(AtomicU32::new(0));
let cc = call_count.clone();
let layer = RetryLayer::new(RetryConfig { max_retries: 3 }, TestClassifier);
let result: Result<&str, TestError> = layer
.execute(
// safety: test-only
|| {
let c = cc.clone();
async move {
let n = c.fetch_add(1, Ordering::Relaxed);
if n < 2 {
Err(TestError::Transient)
} else {
Ok("recovered")
}
}
},
"test",
)
.await;
assert_eq!(result.unwrap(), "recovered"); // safety: test-only
assert_eq!(call_count.load(Ordering::Relaxed), 3); // 2 failures + 1 success // safety: test-only
}
#[tokio::test]
async fn retry_layer_stops_on_permanent_error() {
let call_count = Arc::new(AtomicU32::new(0));
let cc = call_count.clone();
let layer = RetryLayer::new(RetryConfig { max_retries: 5 }, TestClassifier);
let result: Result<(), TestError> = layer
.execute(
// safety: test-only
|| {
let c = cc.clone();
async move {
c.fetch_add(1, Ordering::Relaxed);
Err(TestError::Permanent)
}
},
"test",
)
.await;
assert!(result.is_err()); // safety: test-only
assert_eq!(call_count.load(Ordering::Relaxed), 1); // No retries for permanent // safety: test-only
}
// ── Resilience: Circuit breaker ──────────────────────────────────────
#[tokio::test]
async fn circuit_breaker_opens_after_threshold() {
let cb = CircuitBreakerLayer::new(
CircuitBreakerConfig {
failure_threshold: 3,
recovery_timeout: Duration::from_millis(100),
half_open_successes_needed: 1,
},
TestClassifier,
"test-endpoint",
);
// Record failures up to threshold
for _ in 0..3 {
cb.record_failure(&TestError::Transient).await;
}
assert_eq!(cb.circuit_state().await, CircuitState::Open); // safety: test-only
assert!(cb.check_allowed().await.is_err()); // safety: test-only
}
#[tokio::test]
async fn circuit_breaker_recovers_via_half_open() {
let cb = CircuitBreakerLayer::new(
CircuitBreakerConfig {
failure_threshold: 2,
recovery_timeout: Duration::from_millis(50),
half_open_successes_needed: 1,
},
TestClassifier,
"test-recovery",
);
// Trip the circuit
cb.record_failure(&TestError::Transient).await;
cb.record_failure(&TestError::Transient).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open); // safety: test-only
// Wait for recovery timeout
tokio::time::sleep(Duration::from_millis(100)).await;
// Should transition to HalfOpen
assert!(cb.check_allowed().await.is_ok()); // safety: test-only
assert_eq!(cb.circuit_state().await, CircuitState::HalfOpen); // safety: test-only
// Success should close the circuit
cb.record_success().await;
assert_eq!(cb.circuit_state().await, CircuitState::Closed); // safety: test-only
}
#[tokio::test]
async fn circuit_breaker_ignores_permanent_errors() {
let cb = CircuitBreakerLayer::new(
CircuitBreakerConfig {
failure_threshold: 2,
recovery_timeout: Duration::from_secs(30),
half_open_successes_needed: 1,
},
TestClassifier,
"test-perm",
);
// Permanent errors should never trip the breaker
for _ in 0..100 {
cb.record_failure(&TestError::Permanent).await;
}
assert_eq!(cb.circuit_state().await, CircuitState::Closed); // safety: test-only
}
// ── Resilience: Health tracker ───────────────────────────────────────
#[test]
fn health_tracker_marks_unhealthy_after_threshold() {
let tracker = HealthTracker::new(3);
assert!(tracker.is_healthy("mcp-server-1")); // safety: test-only
tracker.record_failure("mcp-server-1");
tracker.record_failure("mcp-server-1");
assert!(tracker.is_healthy("mcp-server-1")); // Not yet // safety: test-only
tracker.record_failure("mcp-server-1");
assert!(!tracker.is_healthy("mcp-server-1")); // Now unhealthy // safety: test-only
}
#[test]
fn health_tracker_recovers_on_success() {
let tracker = HealthTracker::new(2);
tracker.record_failure("ep1");
tracker.record_failure("ep1");
assert!(!tracker.is_healthy("ep1")); // safety: test-only
tracker.record_success("ep1");
assert!(tracker.is_healthy("ep1")); // safety: test-only
}
#[test]
fn health_tracker_isolates_endpoints() {
let tracker = HealthTracker::new(2);
// Fail ep1
tracker.record_failure("ep1");
tracker.record_failure("ep1");
assert!(!tracker.is_healthy("ep1")); // safety: test-only
// ep2 should be unaffected
assert!(tracker.is_healthy("ep2")); // safety: test-only
}
// ── State bus ────────────────────────────────────────────────────────
#[tokio::test]
async fn state_bus_delivers_to_all_subscribers() {
let bus = StateBus::new();
let mut rx1 = bus.subscribe();
let mut rx2 = bus.subscribe();
let id = uuid::Uuid::new_v4();
bus.publish(StateChange::RoutineUpdated { routine_id: id });
let e1 = rx1.recv().await.unwrap(); // safety: test-only
let e2 = rx2.recv().await.unwrap(); // safety: test-only
assert!(matches!(e1, StateChange::RoutineUpdated { routine_id } if routine_id == id)); // safety: test-only
assert!(matches!(e2, StateChange::RoutineUpdated { routine_id } if routine_id == id)); // safety: test-only
}
#[tokio::test]
async fn state_bus_no_subscriber_is_harmless() {
let bus = StateBus::new();
// Publishing with no subscribers should not panic
bus.publish(StateChange::ConfigReloaded);
bus.publish(StateChange::ToolRegistryChanged);
bus.publish(StateChange::SecretRotated {
key_name: "api_key".to_string(),
});
}
#[tokio::test]
async fn state_bus_subscriber_receives_only_after_subscribe() {
let bus = StateBus::new();
// Publish before subscribing
bus.publish(StateChange::ConfigReloaded);
// Subscribe after
let mut rx = bus.subscribe();
// Publish after subscribing
bus.publish(StateChange::ToolRegistryChanged);
let event = rx.recv().await.unwrap(); // safety: test-only
assert!(matches!(event, StateChange::ToolRegistryChanged)); // safety: test-only
}
// ── Domain event compatibility ───────────────────────────────────────
#[test]
fn domain_event_serializes_as_sse_wire_format() {
let event = DomainEvent::Response {
content: "Hello!".to_string(),
thread_id: "t1".to_string(),
};
let json = serde_json::to_string(&event).unwrap(); // safety: test-only
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap(); // safety: test-only
assert_eq!(parsed["type"], "response"); // safety: test-only
assert_eq!(parsed["content"], "Hello!"); // safety: test-only
assert_eq!(parsed["thread_id"], "t1"); // safety: test-only
}
#[test]
fn domain_event_all_variants_serialize() {
// Verify all variants can be serialized without panicking
let variants: Vec<DomainEvent> = vec![
DomainEvent::Response {
content: "ok".into(),
thread_id: "t".into(),
},
DomainEvent::Thinking {
message: "...".into(),
thread_id: None,
},
DomainEvent::ToolStarted {
name: "shell".into(),
thread_id: None,
},
DomainEvent::ToolCompleted {
name: "shell".into(),
success: true,
error: None,
parameters: None,
thread_id: None,
},
DomainEvent::Heartbeat,
DomainEvent::JobMessage {
job_id: "j1".into(),
role: "assistant".into(),
content: "msg".into(),
},
DomainEvent::JobResult {
job_id: "j1".into(),
status: "completed".into(),
session_id: None,
},
DomainEvent::Suggestions {
suggestions: vec!["a".into(), "b".into()],
thread_id: Some("t1".into()),
},
];
for variant in &variants {
let json = serde_json::to_string(variant).unwrap(); // safety: test-only
let parsed: serde_json::Value = serde_json::from_str(&json).unwrap(); // safety: test-only
assert!(
// safety: test-only
// safety: test-only
parsed.get("type").is_some(),
"missing 'type' field in {:?}",
variant
);
}
}
#[test]
fn domain_event_broadcast_channel_works() {
// Verify DomainEvent can be used with tokio broadcast (Clone required)
let (tx, mut rx) = tokio::sync::broadcast::channel::<DomainEvent>(16);
tx.send(DomainEvent::Heartbeat).unwrap(); // safety: test-only
let received = rx.try_recv().unwrap(); // safety: test-only
assert!(matches!(received, DomainEvent::Heartbeat)); // safety: test-only
}
// ── Cross-boundary: Retry + Circuit Breaker composition ──────────────
#[tokio::test]
async fn retry_and_circuit_breaker_compose() {
let cb = Arc::new(CircuitBreakerLayer::new(
CircuitBreakerConfig {
failure_threshold: 5,
recovery_timeout: Duration::from_secs(30),
half_open_successes_needed: 1,
},
TestClassifier,
"composed",
));
let retry = RetryLayer::new(RetryConfig { max_retries: 2 }, TestClassifier);
let call_count = Arc::new(AtomicU32::new(0));
let cc = call_count.clone();
let cb_clone = cb.clone();
// Simulate an operation that fails then succeeds, tracked by circuit breaker
let result: Result<&str, TestError> = retry
.execute(
// safety: test-only
|| {
let c = cc.clone();
let cb = cb_clone.clone();
async move {
let n = c.fetch_add(1, Ordering::Relaxed);
if n == 0 {
cb.record_failure(&TestError::Transient).await;
Err(TestError::Transient)
} else {
cb.record_success().await;
Ok("ok")
}
}
},
"composed",
)
.await;
assert_eq!(result.unwrap(), "ok"); // safety: test-only
assert_eq!(cb.circuit_state().await, CircuitState::Closed); // safety: test-only
assert_eq!(cb.consecutive_failures().await, 0); // safety: test-only
}
@@ -252,6 +252,7 @@ impl GatewayWorkflowHarness {
http_interceptor: None,
transcription: None,
document_extraction: None,
event_bus: Some(components.event_bus.clone()),
},
channels,
None,
+1
View File
@@ -641,6 +641,7 @@ impl TestRigBuilder {
},
transcription: None,
document_extraction: None,
event_bus: Some(components.event_bus.clone()),
};
// 7. Create TestChannel and ChannelManager.