Merge pull request #1185 from nearai/staging-promote/71b1a677-23096345848

chore: promote staging to staging-promote/579c4fdb-23095333790 (2026-03-14 21:05 UTC)
This commit is contained in:
Henry Park
2026-03-16 07:49:51 -07:00
committed by GitHub
28 changed files with 1932 additions and 106 deletions
Generated
+13 -13
View File
@@ -151,7 +151,7 @@ version = "1.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc"
dependencies = [
"windows-sys 0.61.2",
"windows-sys 0.60.2",
]
[[package]]
@@ -162,7 +162,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d"
dependencies = [
"anstyle",
"once_cell_polyfill",
"windows-sys 0.61.2",
"windows-sys 0.60.2",
]
[[package]]
@@ -2077,7 +2077,7 @@ dependencies = [
"libc",
"option-ext",
"redox_users 0.5.2",
"windows-sys 0.61.2",
"windows-sys 0.59.0",
]
[[package]]
@@ -2264,7 +2264,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -4089,7 +4089,7 @@ version = "0.50.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
dependencies = [
"windows-sys 0.61.2",
"windows-sys 0.59.0",
]
[[package]]
@@ -5433,7 +5433,7 @@ dependencies = [
"errno",
"libc",
"linux-raw-sys 0.12.1",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -6115,7 +6115,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e"
dependencies = [
"libc",
"windows-sys 0.61.2",
"windows-sys 0.60.2",
]
[[package]]
@@ -6337,10 +6337,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.2",
"getrandom 0.3.4",
"once_cell",
"rustix 1.1.4",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@@ -7134,13 +7134,13 @@ checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
[[package]]
name = "uds_windows"
version = "1.2.0"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51b70b87d15e91f553711b40df3048faf27a7a04e01e0ddc0cf9309f0af7c2ca"
checksum = "f2f6fb2847f6742cd76af783a2a2c49e9375d0a111c7bef6f71cd9e738c72d6e"
dependencies = [
"memoffset",
"tempfile",
"windows-sys 0.61.2",
"windows-sys 0.60.2",
]
[[package]]
@@ -7996,7 +7996,7 @@ version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.61.2",
"windows-sys 0.48.0",
]
[[package]]
+1 -1
View File
@@ -2,7 +2,7 @@
"name": "discord",
"display_name": "Discord Channel",
"kind": "channel",
"version": "0.2.1",
"version": "0.2.0",
"wit_version": "0.3.0",
"description": "Talk to your agent in Discord",
"keywords": [
+1 -1
View File
@@ -2,7 +2,7 @@
"name": "github",
"display_name": "GitHub",
"kind": "tool",
"version": "0.2.1",
"version": "0.2.0",
"wit_version": "0.3.0",
"description": "GitHub integration for issues, PRs, repos, and code search",
"keywords": [
+1 -1
View File
@@ -2,7 +2,7 @@
"name": "web-search",
"display_name": "Web Search",
"kind": "tool",
"version": "0.2.1",
"version": "0.2.0",
"wit_version": "0.3.0",
"description": "Search the web using Brave Search API",
"keywords": [
+198 -1
View File
@@ -538,11 +538,174 @@ pub fn next_cron_fire(
}
}
/// Describe common routine cron patterns in plain English.
///
/// Falls back to `cron: <raw>` for malformed or complex expressions.
pub fn describe_cron(schedule: &str, timezone: Option<&str>) -> String {
fn fallback(raw: &str) -> String {
if raw.trim().is_empty() {
"cron: (empty)".to_string()
} else {
format!("cron: {}", raw.trim())
}
}
fn parse_u8_token(token: &str) -> Option<u8> {
token.parse::<u8>().ok()
}
fn parse_step(token: &str) -> Option<u8> {
token
.strip_prefix("*/")
.and_then(parse_u8_token)
.filter(|n| *n > 0)
}
fn weekday_name(dow: &str) -> Option<&'static str> {
let normalized = dow.trim().to_ascii_uppercase();
match normalized.as_str() {
"MON" | "1" => Some("Monday"),
"TUE" | "2" => Some("Tuesday"),
"WED" | "3" => Some("Wednesday"),
"THU" | "4" => Some("Thursday"),
"FRI" | "5" => Some("Friday"),
"SAT" | "6" => Some("Saturday"),
"SUN" | "0" | "7" => Some("Sunday"),
_ => None,
}
}
fn format_time(hour: u8, minute: u8) -> String {
if hour == 0 && minute == 0 {
return "midnight".to_string();
}
let (display_hour, am_pm) = match hour {
0 => (12, "AM"),
1..=11 => (hour, "AM"),
12 => (12, "PM"),
_ => (hour - 12, "PM"),
};
format!("{display_hour}:{minute:02} {am_pm}")
}
fn ordinal(n: u8) -> String {
let suffix = if (11..=13).contains(&(n % 100)) {
"th"
} else {
match n % 10 {
1 => "st",
2 => "nd",
3 => "rd",
_ => "th",
}
};
format!("{n}{suffix}")
}
fn describe_inner(raw: &str) -> Option<String> {
let fields: Vec<&str> = raw.split_whitespace().collect();
let (sec, min, hour, dom, month, dow, year) = match fields.len() {
5 => (
"0", fields[0], fields[1], fields[2], fields[3], fields[4], None,
),
6 => (
fields[0], fields[1], fields[2], fields[3], fields[4], fields[5], None,
),
7 => (
fields[0],
fields[1],
fields[2],
fields[3],
fields[4],
fields[5],
Some(fields[6]),
),
_ => return None,
};
if year.is_some_and(|v| v != "*") {
return None;
}
if sec == "0"
&& hour == "*"
&& dom == "*"
&& month == "*"
&& dow == "*"
&& let Some(step) = parse_step(min)
{
return Some(match step {
1 => "Every minute".to_string(),
n => format!("Every {n} minutes"),
});
}
if sec == "0"
&& min == "0"
&& dom == "*"
&& month == "*"
&& dow == "*"
&& let Some(step) = parse_step(hour)
{
return Some(match step {
1 => "Every hour".to_string(),
n => format!("Every {n} hours"),
});
}
let hour = parse_u8_token(hour).filter(|h| *h <= 23)?;
let minute = parse_u8_token(min).filter(|m| *m <= 59)?;
let time = format_time(hour, minute);
let time_phrase = if time == "midnight" {
"at midnight".to_string()
} else {
format!("at {time}")
};
if sec == "0" && dom == "*" && month == "*" && dow == "*" {
return Some(format!("Daily {time_phrase}"));
}
if sec == "0" && dom == "*" && month == "*" && dow.eq_ignore_ascii_case("MON-FRI") {
return Some(format!("Weekdays {time_phrase}"));
}
if sec == "0"
&& dom == "*"
&& month == "*"
&& let Some(day_name) = weekday_name(dow)
{
return Some(format!("Every {day_name} {time_phrase}"));
}
if sec == "0"
&& month == "*"
&& dow == "*"
&& let Some(day_of_month) = parse_u8_token(dom).filter(|d| (1..=31).contains(d))
{
return Some(format!(
"{} of every month {time_phrase}",
ordinal(day_of_month)
));
}
None
}
let mut description = describe_inner(schedule).unwrap_or_else(|| fallback(schedule));
if let Some(tz) = timezone.map(str::trim).filter(|tz| !tz.is_empty()) {
description.push_str(" (");
description.push_str(tz);
description.push(')');
}
description
}
#[cfg(test)]
mod tests {
use crate::agent::routine::{
MAX_TOOL_ROUNDS_LIMIT, RoutineAction, RoutineGuardrails, RunStatus, Trigger, content_hash,
next_cron_fire,
describe_cron, next_cron_fire,
};
#[test]
@@ -698,6 +861,40 @@ mod tests {
assert_ne!(next_utc, next_est, "timezone should shift the fire time");
}
#[test]
fn test_describe_cron_common_patterns() {
let cases = vec![
("0 */30 * * * *", None, "Every 30 minutes"),
("0 0 9 * * *", None, "Daily at 9:00 AM"),
("0 0 9 * * MON-FRI", None, "Weekdays at 9:00 AM"),
("0 0 */2 * * *", None, "Every 2 hours"),
("0 0 0 * * *", None, "Daily at midnight"),
("0 0 9 * * 1", None, "Every Monday at 9:00 AM"),
("0 0 9 1 * *", None, "1st of every month at 9:00 AM"),
(
"0 0 9 * * MON-FRI",
Some("America/New_York"),
"Weekdays at 9:00 AM (America/New_York)",
),
("1 2 3 4 5 6", None, "cron: 1 2 3 4 5 6"),
];
for (schedule, timezone, expected) in cases {
let actual = describe_cron(schedule, timezone);
assert_eq!(actual, expected); // safety: test-only assertion in #[cfg(test)] module
}
}
#[test]
fn test_describe_cron_edge_cases() {
assert_eq!(describe_cron("", None), "cron: (empty)"); // safety: test-only assertion in #[cfg(test)] module
assert_eq!(describe_cron("not a cron", None), "cron: not a cron"); // safety: test-only assertion in #[cfg(test)] module
let weekdays_5_field = describe_cron("0 9 * * MON-FRI", None);
assert_eq!(weekdays_5_field, "Weekdays at 9:00 AM"); // safety: test-only assertion in #[cfg(test)] module
let weekdays_7_field = describe_cron("0 0 9 * * MON-FRI *", None);
assert_eq!(weekdays_7_field, "Weekdays at 9:00 AM"); // safety: test-only assertion in #[cfg(test)] module
}
#[test]
fn test_guardrails_default() {
let g = RoutineGuardrails::default();
+117 -4
View File
@@ -139,6 +139,32 @@ impl RoutineEngine {
let cache = self.event_cache.read().await;
let mut fired = 0;
// Collect routine IDs for batch query
let routine_ids: Vec<Uuid> = cache
.iter()
.filter_map(|matcher| match matcher {
EventMatcher::Message { routine, .. } => Some(routine.id),
EventMatcher::System { .. } => None,
})
.collect();
if routine_ids.is_empty() {
return 0;
}
// Single batch query instead of N queries
let concurrent_counts = match self
.store
.count_running_routine_runs_batch(&routine_ids)
.await
{
Ok(counts) => counts,
Err(e) => {
tracing::error!("Failed to batch-load concurrent counts: {}", e);
return 0;
}
};
for matcher in cache.iter() {
let (routine, re) = match matcher {
EventMatcher::Message { routine, regex } => (routine, regex),
@@ -164,8 +190,9 @@ impl RoutineEngine {
continue;
}
// Concurrent run check
if !self.check_concurrent(routine).await {
// Concurrent run check (using batch-loaded counts)
let running_count = concurrent_counts.get(&routine.id).copied().unwrap_or(0);
if running_count >= routine.guardrails.max_concurrent as i64 {
tracing::trace!(routine = %routine.name, "Skipped: max concurrent reached");
continue;
}
@@ -197,6 +224,35 @@ impl RoutineEngine {
let cache = self.event_cache.read().await;
let mut fired = 0;
// Collect routine IDs for batch query
let routine_ids: Vec<Uuid> = cache
.iter()
.filter_map(|matcher| match matcher {
EventMatcher::System { routine } => Some(routine.id),
EventMatcher::Message { .. } => None,
})
.collect();
if routine_ids.is_empty() {
return 0;
}
// Single batch query instead of N queries
let concurrent_counts = match self
.store
.count_running_routine_runs_batch(&routine_ids)
.await
{
Ok(counts) => counts,
Err(e) => {
tracing::error!(
"Failed to batch-load concurrent counts for system events: {}",
e
);
return 0;
}
};
for matcher in cache.iter() {
let routine = match matcher {
EventMatcher::System { routine } => routine,
@@ -248,7 +304,9 @@ impl RoutineEngine {
continue;
}
if !self.check_concurrent(routine).await {
// Concurrent run check (using batch-loaded counts)
let running_count = concurrent_counts.get(&routine.id).copied().unwrap_or(0);
if running_count >= routine.guardrails.max_concurrent as i64 {
tracing::debug!(routine = %routine.name, "Skipped: max concurrent reached");
continue;
}
@@ -925,7 +983,8 @@ async fn execute_lightweight_with_tools(
.tool_definitions_excluding(ROUTINE_TOOL_DENYLIST)
.await;
let request = ToolCompletionRequest::new(messages.clone(), tool_defs)
let request_messages = snapshot_messages_for_tool_iteration(&messages);
let request = ToolCompletionRequest::new(request_messages, tool_defs)
.with_max_tokens(effective_max_tokens)
.with_temperature(0.3);
@@ -1001,6 +1060,31 @@ async fn execute_lightweight_with_tools(
}
}
// Bound per-iteration context copy cost for lightweight tool loops.
const MAX_TOOL_LOOP_MESSAGES: usize = 32;
fn snapshot_messages_for_tool_iteration(messages: &[ChatMessage]) -> Vec<ChatMessage> {
if messages.len() <= MAX_TOOL_LOOP_MESSAGES {
return messages.to_vec();
}
let mut snapshot = Vec::with_capacity(MAX_TOOL_LOOP_MESSAGES);
if let Some(first) = messages.first()
&& first.role == crate::llm::Role::System
{
snapshot.push(first.clone());
let tail_len = MAX_TOOL_LOOP_MESSAGES - 1;
let tail_start = (messages.len() - tail_len).max(1);
snapshot.extend_from_slice(&messages[tail_start..]);
} else {
let tail_start = messages.len() - MAX_TOOL_LOOP_MESSAGES;
snapshot.extend_from_slice(&messages[tail_start..]);
}
snapshot
}
/// Tools that must never be callable from lightweight routines.
///
/// These tools pose autonomy-escalation risks: a routine could self-replicate,
@@ -1386,4 +1470,33 @@ mod tests {
let out = super::truncate(input, 5);
assert_eq!(out, "abcde...");
}
#[test]
fn test_snapshot_messages_keeps_system_and_recent_tail() {
let mut messages = vec![crate::llm::ChatMessage::system("sys")];
for i in 0..80 {
messages.push(crate::llm::ChatMessage::user(format!("u{i}")));
}
let snapshot = super::snapshot_messages_for_tool_iteration(&messages);
assert_eq!(snapshot.len(), super::MAX_TOOL_LOOP_MESSAGES); // safety: test-only no-panics CI false positive
assert_eq!(snapshot[0].role, crate::llm::Role::System); // safety: test-only no-panics CI false positive
assert_eq!(snapshot[0].content, "sys"); // safety: test-only no-panics CI false positive
let last_content = snapshot.last().map(|m| m.content.as_str());
assert_eq!(last_content, Some("u79")); // safety: test-only no-panics CI false positive
}
#[test]
fn test_snapshot_messages_unchanged_when_within_limit() {
let messages = vec![
crate::llm::ChatMessage::system("sys"),
crate::llm::ChatMessage::user("a"),
crate::llm::ChatMessage::assistant("b"),
];
let snapshot = super::snapshot_messages_for_tool_iteration(&messages);
assert_eq!(snapshot.len(), messages.len()); // safety: test-only no-panics CI false positive
assert_eq!(snapshot[0].role, crate::llm::Role::System); // safety: test-only no-panics CI false positive
assert_eq!(snapshot[1].content, "a"); // safety: test-only no-panics CI false positive
assert_eq!(snapshot[2].content, "b"); // safety: test-only no-panics CI false positive
}
}
+33 -9
View File
@@ -179,27 +179,33 @@ impl Scheduler {
})
.unwrap_or(self.config.max_tokens_per_job);
// Apply both metadata and token budget in one closure (Issue #813: atomic update)
if let Some(meta) = metadata {
// Apply both metadata and token budget in one closure (Issue #813: atomic update).
// Use update_context_and_get to ensure atomicity: no gap where concurrent workers
// can modify the context between update and DB persist (Issue #807).
let ctx = if let Some(meta) = metadata {
self.context_manager
.update_context(job_id, |ctx| {
.update_context_and_get(job_id, |ctx| {
ctx.metadata = meta;
if max_tokens > 0 {
ctx.max_tokens = max_tokens;
}
})
.await?;
.await?
} else if max_tokens > 0 {
self.context_manager
.update_context(job_id, |ctx| {
.update_context_and_get(job_id, |ctx| {
ctx.max_tokens = max_tokens;
})
.await?;
}
.await?
} else {
// No metadata or token budget to set; get the initial context
self.context_manager.get_context(job_id).await?
};
// Persist to DB before scheduling so the worker's FK references are valid
// Persist to DB before scheduling so the worker's FK references are valid.
// The context was read under the same lock as the update (atomic), preventing
// concurrent worker interference (Issue #807: non-transactional context updates).
if let Some(ref store) = self.store {
let ctx = self.context_manager.get_context(job_id).await?;
store.save_job(&ctx).await.map_err(|e| JobError::Failed {
id: job_id,
reason: format!("failed to persist job: {e}"),
@@ -832,6 +838,24 @@ mod tests {
);
}
#[tokio::test]
async fn test_dispatch_job_no_metadata_no_user_tokens_edge_case() {
// Edge case coverage: when metadata=None AND max_tokens=0 (config),
// the else branch calls get_context() directly (not update_context_and_get).
// This test verifies that path works correctly (Issue #807: full branch coverage).
let sched = make_test_scheduler(0); // 0 = unlimited, but user provides None
let job_id = sched
.dispatch_job("user1", "test", "desc", None) // None metadata
.await
.unwrap(); // safety: test code
let ctx = sched.context_manager.get_context(job_id).await.unwrap(); // safety: test code
// No metadata was set, should have default empty metadata
assert!(ctx.metadata.is_null() || ctx.metadata == serde_json::json!({})); // safety: test code
// No user tokens AND unlimited config means max_tokens stays at default
assert_eq!(ctx.max_tokens, 0, "unlimited config"); // safety: test code
}
#[test]
fn test_scheduler_creation() {
// Would need to mock dependencies for proper testing
+4
View File
@@ -112,12 +112,16 @@ pub async fn routines_detail_handler(
job_id: run.job_id,
})
.collect();
let routine_info = RoutineInfo::from_routine(&routine);
Ok(Json(RoutineDetailResponse {
id: routine.id,
name: routine.name.clone(),
description: routine.description.clone(),
enabled: routine.enabled,
trigger_type: routine_info.trigger_type,
trigger_raw: routine_info.trigger_raw,
trigger_summary: routine_info.trigger_summary,
trigger: serde_json::to_value(&routine.trigger).unwrap_or_default(),
action: serde_json::to_value(&routine.action).unwrap_or_default(),
guardrails: serde_json::to_value(&routine.guardrails).unwrap_or_default(),
+42 -46
View File
@@ -419,6 +419,44 @@ fn parse_stop(val: &serde_json::Value) -> Option<Vec<String>> {
}
}
fn build_completion_request(
req: &OpenAiChatRequest,
messages: Vec<ChatMessage>,
) -> CompletionRequest {
let mut comp_req = CompletionRequest::new(messages).with_model(req.model.clone());
if let Some(t) = req.temperature {
comp_req = comp_req.with_temperature(t);
}
if let Some(mt) = req.max_tokens {
comp_req = comp_req.with_max_tokens(mt);
}
if let Some(stops) = req.stop.as_ref().and_then(parse_stop) {
comp_req.stop_sequences = Some(stops);
}
comp_req
}
fn build_tool_request(
req: &OpenAiChatRequest,
messages: Vec<ChatMessage>,
) -> ToolCompletionRequest {
let tools = convert_tools(req.tools.as_deref().unwrap_or(&[]));
let mut tool_req = ToolCompletionRequest::new(messages, tools).with_model(req.model.clone());
if let Some(t) = req.temperature {
tool_req = tool_req.with_temperature(t);
}
if let Some(mt) = req.max_tokens {
tool_req = tool_req.with_max_tokens(mt);
}
if let Some(stops) = req.stop.as_ref().and_then(parse_stop) {
tool_req = tool_req.with_stop_sequences(stops);
}
if let Some(choice) = req.tool_choice.as_ref().and_then(normalize_tool_choice) {
tool_req = tool_req.with_tool_choice(choice);
}
tool_req
}
// ---------------------------------------------------------------------------
// Handlers
// ---------------------------------------------------------------------------
@@ -476,19 +514,7 @@ pub async fn chat_completions_handler(
let created = unix_timestamp();
if has_tools {
let tools = convert_tools(req.tools.as_deref().unwrap_or(&[]));
let mut tool_req = ToolCompletionRequest::new(messages, tools).with_model(req.model);
if let Some(t) = req.temperature {
tool_req = tool_req.with_temperature(t);
}
if let Some(mt) = req.max_tokens {
tool_req = tool_req.with_max_tokens(mt);
}
if let Some(ref tc) = req.tool_choice
&& let Some(choice) = normalize_tool_choice(tc)
{
tool_req = tool_req.with_tool_choice(choice);
}
let tool_req = build_tool_request(&req, messages);
let resp = llm
.complete_with_tools(tool_req)
@@ -527,16 +553,7 @@ pub async fn chat_completions_handler(
Ok(Json(response).into_response())
} else {
let mut comp_req = CompletionRequest::new(messages).with_model(req.model);
if let Some(t) = req.temperature {
comp_req = comp_req.with_temperature(t);
}
if let Some(mt) = req.max_tokens {
comp_req = comp_req.with_max_tokens(mt);
}
if let Some(ref stop_val) = req.stop {
comp_req.stop_sequences = parse_stop(stop_val);
}
let comp_req = build_completion_request(&req, messages);
let resp = llm.complete(comp_req).await.map_err(map_llm_error)?;
let model_name = llm.effective_model_name(Some(requested_model.as_str()));
@@ -596,35 +613,14 @@ async fn handle_streaming(
}
let llm_result = if has_tools {
let tools = convert_tools(req.tools.as_deref().unwrap_or(&[]));
let mut tool_req = ToolCompletionRequest::new(messages, tools).with_model(req.model);
if let Some(t) = req.temperature {
tool_req = tool_req.with_temperature(t);
}
if let Some(mt) = req.max_tokens {
tool_req = tool_req.with_max_tokens(mt);
}
if let Some(ref tc) = req.tool_choice
&& let Some(choice) = normalize_tool_choice(tc)
{
tool_req = tool_req.with_tool_choice(choice);
}
let tool_req = build_tool_request(&req, messages);
LlmResult::WithTools(
llm.complete_with_tools(tool_req)
.await
.map_err(map_llm_error)?,
)
} else {
let mut comp_req = CompletionRequest::new(messages).with_model(req.model);
if let Some(t) = req.temperature {
comp_req = comp_req.with_temperature(t);
}
if let Some(mt) = req.max_tokens {
comp_req = comp_req.with_max_tokens(mt);
}
if let Some(ref stop_val) = req.stop {
comp_req.stop_sequences = parse_stop(stop_val);
}
let comp_req = build_completion_request(&req, messages);
LlmResult::Simple(llm.complete(comp_req).await.map_err(map_llm_error)?)
};
let model_name = llm.effective_model_name(Some(requested_model.as_str()));
+4
View File
@@ -2346,12 +2346,16 @@ async fn routines_detail_handler(
job_id: run.job_id,
})
.collect();
let routine_info = RoutineInfo::from_routine(&routine);
Ok(Json(RoutineDetailResponse {
id: routine.id,
name: routine.name.clone(),
description: routine.description.clone(),
enabled: routine.enabled,
trigger_type: routine_info.trigger_type,
trigger_raw: routine_info.trigger_raw,
trigger_summary: routine_info.trigger_summary,
trigger: serde_json::to_value(&routine.trigger).unwrap_or_default(),
action: serde_json::to_value(&routine.action).unwrap_or_default(),
guardrails: serde_json::to_value(&routine.guardrails).unwrap_or_default(),
+21 -3
View File
@@ -3535,10 +3535,13 @@ function renderRoutinesList(routines) {
const toggleLabel = r.enabled ? 'Disable' : 'Enable';
const toggleClass = r.enabled ? 'btn-cancel' : 'btn-restart';
const triggerTitle = (r.trigger_type === 'cron' && r.trigger_raw)
? ' title="' + escapeHtml(r.trigger_raw) + '"'
: '';
return '<tr class="routine-row" data-action="open-routine" data-id="' + escapeHtml(r.id) + '">'
+ '<td>' + escapeHtml(r.name) + '</td>'
+ '<td>' + escapeHtml(r.trigger_summary) + '</td>'
+ '<td' + triggerTitle + '>' + escapeHtml(r.trigger_summary) + '</td>'
+ '<td>' + escapeHtml(r.action_type) + '</td>'
+ '<td>' + formatRelativeTime(r.last_run_at) + '</td>'
+ '<td>' + formatRelativeTime(r.next_fire_at) + '</td>'
@@ -3606,8 +3609,23 @@ function renderRoutineDetail(routine) {
}
// Trigger config
html += '<div class="job-description"><h3>Trigger</h3>'
+ '<pre class="action-json">' + escapeHtml(JSON.stringify(routine.trigger, null, 2)) + '</pre></div>';
if (routine.trigger_type === 'cron') {
const summary = routine.trigger_summary || 'cron';
const raw = routine.trigger_raw || '';
const timezone = routine.trigger && routine.trigger.timezone ? String(routine.trigger.timezone) : '';
html += '<div class="job-description"><h3>Trigger</h3>'
+ '<div class="job-description-body"><strong>' + escapeHtml(summary) + '</strong></div>';
if (raw) {
html += '<div class="job-meta-item">'
+ '<span class="job-meta-label">Raw</span>'
+ '<span class="job-meta-value">' + escapeHtml(raw + (timezone ? ' (' + timezone + ')' : '')) + '</span>'
+ '</div>';
}
html += '</div>';
} else {
html += '<div class="job-description"><h3>Trigger</h3>'
+ '<pre class="action-json">' + escapeHtml(JSON.stringify(routine.trigger, null, 2)) + '</pre></div>';
}
// Action config
html += '<div class="job-description"><h3>Action</h3>'
+22 -8
View File
@@ -735,6 +735,7 @@ pub struct RoutineInfo {
pub description: String,
pub enabled: bool,
pub trigger_type: String,
pub trigger_raw: String,
pub trigger_summary: String,
pub action_type: String,
pub last_run_at: Option<String>,
@@ -747,25 +748,34 @@ pub struct RoutineInfo {
impl RoutineInfo {
/// Convert a `Routine` to the trimmed `RoutineInfo` for list display.
pub fn from_routine(r: &crate::agent::routine::Routine) -> Self {
let (trigger_type, trigger_summary) = match &r.trigger {
crate::agent::routine::Trigger::Cron { schedule, .. } => {
("cron".to_string(), format!("cron: {}", schedule))
}
let (trigger_type, trigger_raw, trigger_summary) = match &r.trigger {
crate::agent::routine::Trigger::Cron { schedule, timezone } => (
"cron".to_string(),
schedule.clone(),
crate::agent::routine::describe_cron(schedule, timezone.as_deref()),
),
crate::agent::routine::Trigger::Event {
pattern, channel, ..
} => {
let ch = channel.as_deref().unwrap_or("any");
("event".to_string(), format!("on {} /{}/", ch, pattern))
(
"event".to_string(),
String::new(),
format!("on {} /{}/", ch, pattern),
)
}
crate::agent::routine::Trigger::SystemEvent {
source, event_type, ..
} => (
"system_event".to_string(),
String::new(),
format!("event: {}.{}", source, event_type),
),
crate::agent::routine::Trigger::Manual => {
("manual".to_string(), "manual only".to_string())
}
crate::agent::routine::Trigger::Manual => (
"manual".to_string(),
String::new(),
"manual only".to_string(),
),
};
let action_type = match &r.action {
@@ -787,6 +797,7 @@ impl RoutineInfo {
description: r.description.clone(),
enabled: r.enabled,
trigger_type,
trigger_raw,
trigger_summary,
action_type: action_type.to_string(),
last_run_at: r.last_run_at.map(|dt| dt.to_rfc3339()),
@@ -818,6 +829,9 @@ pub struct RoutineDetailResponse {
pub name: String,
pub description: String,
pub enabled: bool,
pub trigger_type: String,
pub trigger_raw: String,
pub trigger_summary: String,
pub trigger: serde_json::Value,
pub action: serde_json::Value,
pub guardrails: serde_json::Value,
+38 -2
View File
@@ -139,12 +139,19 @@ impl WebhookServer {
self.config.addr
}
/// Take ownership of shutdown primitives so callers can perform async
/// shutdown work without holding external locks around this server.
pub fn begin_shutdown(&mut self) -> (Option<oneshot::Sender<()>>, Option<JoinHandle<()>>) {
(self.shutdown_tx.take(), self.handle.take())
}
/// Signal graceful shutdown and wait for the server task to finish.
pub async fn shutdown(&mut self) {
if let Some(tx) = self.shutdown_tx.take() {
let (shutdown_tx, handle) = self.begin_shutdown();
if let Some(tx) = shutdown_tx {
let _ = tx.send(());
}
if let Some(handle) = self.handle.take() {
if let Some(handle) = handle {
let _ = handle.await;
}
}
@@ -269,6 +276,35 @@ mod tests {
server.shutdown().await;
}
#[tokio::test]
async fn test_begin_shutdown_takes_handles_for_lock_free_shutdown() {
let addr = SocketAddr::from((std::net::Ipv4Addr::LOCALHOST, 0));
let mut server = WebhookServer::new(WebhookServerConfig { addr });
let test_router = axum::Router::new().route(
"/health",
axum::routing::get(|| async { Json(json!({"status": "ok"})) }),
);
server.add_routes(test_router);
server.start().await.expect("Failed to start server"); // safety: test assertion for setup precondition
let (shutdown_tx, handle) = server.begin_shutdown();
assert!(shutdown_tx.is_some(), "shutdown sender should be available"); // safety: test assertion for expected server state
assert!(handle.is_some(), "server handle should be available"); // safety: test assertion for expected server state
// begin_shutdown() should leave no handles behind on the server.
let (shutdown_tx2, handle2) = server.begin_shutdown();
assert!(shutdown_tx2.is_none(), "shutdown sender should be consumed"); // safety: test assertion for postcondition
assert!(handle2.is_none(), "server handle should be consumed"); // safety: test assertion for postcondition
if let Some(tx) = shutdown_tx {
let _ = tx.send(());
}
if let Some(handle) = handle {
let _ = handle.await;
}
}
#[tokio::test]
async fn test_restart_with_addr_rollback_on_bind_failure() {
use std::net::TcpListener as StdTcpListener;
+88
View File
@@ -87,6 +87,28 @@ impl ContextManager {
Ok(f(context))
}
/// Atomically update a job context and return the updated context.
///
/// This method holds the write lock for the entire update-and-read sequence,
/// preventing concurrent workers from interleaving modifications between the
/// update and the subsequent read (Issue #807: non-transactional context updates).
/// Use this when you need to update context and immediately persist it to DB.
pub async fn update_context_and_get<F>(
&self,
job_id: Uuid,
f: F,
) -> Result<JobContext, JobError>
where
F: FnOnce(&mut JobContext),
{
let mut contexts = self.contexts.write().await;
let context = contexts
.get_mut(&job_id)
.ok_or(JobError::NotFound { id: job_id })?;
f(context);
Ok(context.clone())
}
/// Get job memory.
pub async fn get_memory(&self, job_id: Uuid) -> Result<Memory, JobError> {
self.memories
@@ -877,4 +899,70 @@ mod tests {
assert_eq!(manager.all_jobs().await.len(), 10);
}
#[tokio::test]
async fn update_context_and_get_atomicity_regression_issue_807() {
// Regression test for Issue #807: non-transactional context updates.
// Verify that update_context_and_get returns the exact state that was set,
// without allowing concurrent workers to interleave modifications.
let manager = std::sync::Arc::new(ContextManager::new(100));
let job_id = manager
.create_job("Atomicity Test", "verify no race condition")
.await
.unwrap(); // safety: test code
// Update and get atomically, setting metadata
let metadata = serde_json::json!({ "priority": "high", "user_id": 42 });
let returned_ctx = manager
.update_context_and_get(job_id, |ctx| {
ctx.metadata = metadata.clone();
ctx.max_tokens = 5000;
})
.await
.unwrap(); // safety: test code
// Verify the returned context has the exact updates we set
assert_eq!(returned_ctx.metadata, metadata); // safety: test code
assert_eq!(returned_ctx.max_tokens, 5000); // safety: test code
// Verify a fresh get returns the same state
let fresh_ctx = manager.get_context(job_id).await.unwrap(); // safety: test code
assert_eq!(fresh_ctx.metadata, metadata); // safety: test code
assert_eq!(fresh_ctx.max_tokens, 5000); // safety: test code
}
#[tokio::test]
async fn update_context_and_get_no_concurrent_interleave() {
// Verify that concurrent updates cannot interleave during update_context_and_get.
// If the lock were released too early, a concurrent state transition could
// get mixed into the returned context.
let manager = std::sync::Arc::new(ContextManager::new(100));
let job_id = manager
.create_job("Concurrent Race Test", "ensure atomicity")
.await
.unwrap(); // safety: test code
let metadata = serde_json::json!({ "test": "race_condition" });
let metadata_clone = metadata.clone();
// Spawn a task that will update_context_and_get
let mgr1 = std::sync::Arc::clone(&manager);
let returned_ctx_handle = tokio::spawn(async move {
mgr1.update_context_and_get(job_id, |ctx| {
ctx.metadata = metadata_clone;
ctx.max_tokens = 3000;
})
.await
});
// The returned context should have *only* the metadata update, not any
// concurrent state transitions that might happen during the operation.
let returned_ctx = returned_ctx_handle.await.unwrap().unwrap(); // safety: test code
// Verify atomicity: returned context has the metadata we set
assert_eq!(returned_ctx.metadata, metadata); // safety: test code
assert_eq!(returned_ctx.max_tokens, 3000); // safety: test code
// And it's in the initial state (Pending), not modified by concurrent workers
assert_eq!(returned_ctx.state, crate::context::JobState::Pending); // safety: test code
}
}
+55 -2
View File
@@ -1,13 +1,15 @@
//! Routine-related RoutineStore implementation for LibSqlBackend.
use std::collections::{HashMap, HashSet};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use libsql::params;
use uuid::Uuid;
use super::{
LibSqlBackend, ROUTINE_COLUMNS, ROUTINE_RUN_COLUMNS, fmt_opt_ts, fmt_ts, get_i64, opt_text,
opt_text_owned, row_to_routine_libsql, row_to_routine_run_libsql,
LibSqlBackend, ROUTINE_COLUMNS, ROUTINE_RUN_COLUMNS, fmt_opt_ts, fmt_ts, get_i64, get_text,
opt_text, opt_text_owned, row_to_routine_libsql, row_to_routine_run_libsql,
};
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
use crate::db::RoutineStore;
@@ -409,6 +411,57 @@ impl RoutineStore for LibSqlBackend {
}
}
async fn count_running_routine_runs_batch(
&self,
routine_ids: &[Uuid],
) -> Result<HashMap<Uuid, i64>, DatabaseError> {
if routine_ids.is_empty() {
return Ok(HashMap::new());
}
let mut counts = HashMap::new();
let conn = self.connect().await?;
// Query all running routines and filter in memory
// This is simpler for libSQL than building dynamic parameter lists
let mut rows = conn
.query(
"SELECT routine_id, COUNT(*) as cnt FROM routine_runs
WHERE status = 'running'
GROUP BY routine_id",
params![],
)
.await
.map_err(|e| {
DatabaseError::Query(format!("Failed to batch count running routines: {}", e))
})?;
let routine_id_set: HashSet<Uuid> = routine_ids.iter().copied().collect();
while let Some(row) = rows
.next()
.await
.map_err(|e| DatabaseError::Query(e.to_string()))?
{
let id_str: String = get_text(&row, 0);
let id = Uuid::parse_str(&id_str)
.map_err(|e| DatabaseError::Query(format!("Invalid routine UUID: {}", e)))?;
// Only include if this routine ID was requested
if routine_id_set.contains(&id) {
let cnt: i64 = get_i64(&row, 1);
counts.insert(id, cnt);
}
}
// Ensure all requested IDs are in the map (defaults to 0 for no running runs)
for id in routine_ids {
counts.entry(*id).or_insert(0);
}
Ok(counts)
}
async fn link_routine_run_to_job(
&self,
run_id: Uuid,
+4
View File
@@ -387,6 +387,10 @@ pub trait RoutineStore: Send + Sync {
limit: i64,
) -> Result<Vec<RoutineRun>, DatabaseError>;
async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError>;
async fn count_running_routine_runs_batch(
&self,
routine_ids: &[Uuid],
) -> Result<HashMap<Uuid, i64>, DatabaseError>;
async fn link_routine_run_to_job(
&self,
run_id: Uuid,
+9
View File
@@ -487,6 +487,15 @@ impl RoutineStore for PgBackend {
self.store.count_running_routine_runs(routine_id).await
}
async fn count_running_routine_runs_batch(
&self,
routine_ids: &[Uuid],
) -> Result<std::collections::HashMap<Uuid, i64>, DatabaseError> {
self.store
.count_running_routine_runs_batch(routine_ids)
.await
}
async fn link_routine_run_to_job(
&self,
run_id: Uuid,
+150 -9
View File
@@ -304,6 +304,34 @@ impl ExtensionManager {
*self.relay_channel_manager.write().await = Some(channel_manager);
}
async fn current_channel_owner_id(&self, name: &str) -> Option<i64> {
{
let rt_guard = self.channel_runtime.read().await;
if let Some(owner_id) = rt_guard
.as_ref()
.and_then(|rt| rt.wasm_channel_owner_ids.get(name).copied())
{
return Some(owner_id);
}
}
let store = self.store.as_ref()?;
let key = format!("channels.wasm_channel_owner_ids.{name}");
match store.get_setting(&self.user_id, &key).await {
Ok(Some(serde_json::Value::Number(n))) => n.as_i64(),
Ok(Some(serde_json::Value::String(s))) => s.parse::<i64>().ok(),
Ok(Some(_)) | Ok(None) => None,
Err(e) => {
tracing::debug!(
channel = %name,
error = %e,
"Failed to read persisted wasm channel owner id"
);
None
}
}
}
/// Check if a channel name corresponds to a relay extension (has stored stream token).
pub async fn is_relay_channel(&self, name: &str) -> bool {
self.secrets
@@ -2980,13 +3008,7 @@ impl ExtensionManager {
// Verify runtime infrastructure is available and clone Arcs so we don't
// hold the RwLock guard across awaits.
let (
channel_runtime,
channel_manager,
pairing_store,
wasm_channel_router,
wasm_channel_owner_ids,
) = {
let (channel_runtime, channel_manager, pairing_store, wasm_channel_router) = {
let rt_guard = self.channel_runtime.read().await;
let rt = rt_guard.as_ref().ok_or_else(|| {
ExtensionError::ActivationFailed("WASM channel runtime not configured".to_string())
@@ -2996,7 +3018,6 @@ impl ExtensionManager {
Arc::clone(&rt.channel_manager),
Arc::clone(&rt.pairing_store),
Arc::clone(&rt.wasm_channel_router),
rt.wasm_channel_owner_ids.clone(),
)
};
@@ -3067,7 +3088,7 @@ impl ExtensionManager {
);
}
if let Some(&owner_id) = wasm_channel_owner_ids.get(channel_name.as_str()) {
if let Some(owner_id) = self.current_channel_owner_id(&channel_name).await {
config_updates.insert("owner_id".to_string(), serde_json::json!(owner_id));
}
@@ -4744,6 +4765,126 @@ mod tests {
)
}
#[tokio::test]
async fn test_current_channel_owner_id_uses_runtime_state() -> Result<(), String> {
let manager = make_manager_with_temp_dirs();
if manager.current_channel_owner_id("telegram").await.is_some() {
return Err("expected no owner id for telegram before runtime setup".to_string());
}
let channels = Arc::new(crate::channels::ChannelManager::new());
let runtime = Arc::new(
crate::channels::wasm::WasmChannelRuntime::new(
crate::channels::wasm::WasmChannelRuntimeConfig::default(),
)
.map_err(|e| format!("runtime init failed: {e}"))?,
);
let pairing_store = Arc::new(crate::pairing::PairingStore::new());
let router = Arc::new(crate::channels::wasm::WasmChannelRouter::new());
let mut owner_ids = std::collections::HashMap::new();
owner_ids.insert("telegram".to_string(), 12345_i64);
manager
.set_channel_runtime(channels, runtime, pairing_store, router, owner_ids)
.await;
if manager.current_channel_owner_id("telegram").await != Some(12345_i64) {
return Err("expected runtime owner id fast-path for telegram".to_string());
}
if manager.current_channel_owner_id("slack").await.is_some() {
return Err("expected no owner id for slack".to_string());
}
Ok(())
}
#[cfg(feature = "libsql")]
#[tokio::test]
async fn test_current_channel_owner_id_uses_store_fallback() -> Result<(), String> {
use crate::db::{Database, SettingsStore};
let dir = tempfile::tempdir().map_err(|e| format!("tempdir failed: {e}"))?;
let db_path = dir.path().join("owner-id.db");
let db = Arc::new(
crate::db::libsql::LibSqlBackend::new_local(&db_path)
.await
.map_err(|e| format!("create local libsql backend failed: {e}"))?,
);
db.run_migrations()
.await
.map_err(|e| format!("run libsql migrations failed: {e}"))?;
let tools_dir = dir.path().join("tools");
let channels_dir = dir.path().join("channels");
std::fs::create_dir_all(&tools_dir).ok();
std::fs::create_dir_all(&channels_dir).ok();
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
use crate::testing::credentials::TEST_CRYPTO_KEY;
use crate::tools::ToolRegistry;
use crate::tools::mcp::process::McpProcessManager;
use crate::tools::mcp::session::McpSessionManager;
let master_key = secrecy::SecretString::from(TEST_CRYPTO_KEY.to_string());
let crypto = Arc::new(
SecretsCrypto::new(master_key)
.map_err(|e| format!("create secrets crypto failed: {e}"))?,
);
let manager = ExtensionManager::new(
Arc::new(McpSessionManager::new()),
Arc::new(McpProcessManager::new()),
Arc::new(InMemorySecretsStore::new(crypto)),
Arc::new(ToolRegistry::new()),
None,
None,
tools_dir,
channels_dir,
None,
"test".to_string(),
Some(db.clone() as Arc<dyn crate::db::Database>),
Vec::new(),
);
if manager.current_channel_owner_id("telegram").await.is_some() {
return Err("expected no owner id before settings seed".to_string());
}
db.set_setting(
"test",
"channels.wasm_channel_owner_ids.telegram",
&serde_json::json!(54321_i64),
)
.await
.map_err(|e| format!("persist owner id in settings failed: {e}"))?;
if manager.current_channel_owner_id("telegram").await != Some(54321_i64) {
return Err("expected store fallback owner id for telegram".to_string());
}
let channels = Arc::new(crate::channels::ChannelManager::new());
let runtime = Arc::new(
crate::channels::wasm::WasmChannelRuntime::new(
crate::channels::wasm::WasmChannelRuntimeConfig::default(),
)
.map_err(|e| format!("runtime init failed: {e}"))?,
);
let pairing_store = Arc::new(crate::pairing::PairingStore::new());
let router = Arc::new(crate::channels::wasm::WasmChannelRouter::new());
let mut owner_ids = std::collections::HashMap::new();
owner_ids.insert("telegram".to_string(), 12345_i64);
manager
.set_channel_runtime(channels, runtime, pairing_store, router, owner_ids)
.await;
if manager.current_channel_owner_id("telegram").await != Some(12345_i64) {
return Err("expected runtime fast-path owner id precedence".to_string());
}
Ok(())
}
// ── resolve_env_credentials tests ────────────────────────────────────
#[test]
+39
View File
@@ -1,5 +1,8 @@
//! PostgreSQL store for persisting agent data.
#[cfg(feature = "postgres")]
use std::collections::HashMap;
use chrono::{DateTime, Utc};
#[cfg(feature = "postgres")]
use deadpool_postgres::{Config, Pool};
@@ -1294,6 +1297,42 @@ impl Store {
Ok(row.get("cnt"))
}
/// Batch-load concurrent run counts for multiple routines in a single query.
/// Returns a map where missing routine IDs default to 0.
#[cfg(feature = "postgres")]
pub async fn count_running_routine_runs_batch(
&self,
routine_ids: &[Uuid],
) -> Result<HashMap<Uuid, i64>, DatabaseError> {
if routine_ids.is_empty() {
return Ok(HashMap::new());
}
let conn = self.conn().await?;
let rows = conn
.query(
"SELECT routine_id, COUNT(*) as cnt FROM routine_runs
WHERE routine_id = ANY($1) AND status = 'running'
GROUP BY routine_id",
&[&routine_ids],
)
.await?;
let mut counts = HashMap::new();
for row in rows {
let id: Uuid = row.get("routine_id");
let cnt: i64 = row.get("cnt");
counts.insert(id, cnt);
}
// Ensure all requested IDs are in the map (defaults to 0 for no running runs)
for id in routine_ids {
counts.entry(*id).or_insert(0);
}
Ok(counts)
}
/// Link a routine run to a dispatched job.
pub async fn link_routine_run_to_job(
&self,
+5 -2
View File
@@ -176,8 +176,11 @@ impl LlmProvider for BedrockProvider {
builder = builder.tool_config(tc);
}
if let Some(config) = build_inference_config(request.temperature, request.max_tokens, None)
{
if let Some(config) = build_inference_config(
request.temperature,
request.max_tokens,
request.stop_sequences.as_deref(),
) {
builder = builder.inference_config(config);
}
+6
View File
@@ -475,6 +475,7 @@ impl LlmProvider for NearAiChatProvider {
messages,
temperature: req.temperature,
max_tokens: req.max_tokens,
stop: req.stop_sequences,
tools: None,
tool_choice: None,
};
@@ -554,6 +555,7 @@ impl LlmProvider for NearAiChatProvider {
messages,
temperature: req.temperature,
max_tokens: req.max_tokens,
stop: req.stop_sequences,
tools: if tools.is_empty() { None } else { Some(tools) },
tool_choice: req.tool_choice,
};
@@ -680,6 +682,8 @@ struct ChatCompletionRequest {
#[serde(skip_serializing_if = "Option::is_none")]
max_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
stop: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
tools: Option<Vec<ChatCompletionTool>>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_choice: Option<String>,
@@ -1666,6 +1670,7 @@ mod tests {
}],
temperature: None,
max_tokens: None,
stop: None,
tools: None,
tool_choice: None,
};
@@ -1687,6 +1692,7 @@ mod tests {
messages: vec![],
temperature: Some(0.7),
max_tokens: Some(1024),
stop: None,
tools: Some(vec![ChatCompletionTool {
tool_type: "function".to_string(),
function: ChatCompletionFunction {
+24 -3
View File
@@ -251,6 +251,7 @@ pub struct ToolCompletionRequest {
pub model: Option<String>,
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
pub stop_sequences: Option<Vec<String>>,
/// How to handle tool use: "auto", "required", or "none".
pub tool_choice: Option<String>,
/// Opaque metadata passed through to the provider (e.g. thread_id for chaining).
@@ -266,6 +267,7 @@ impl ToolCompletionRequest {
model: None,
max_tokens: None,
temperature: None,
stop_sequences: None,
tool_choice: None,
metadata: std::collections::HashMap::new(),
}
@@ -289,6 +291,12 @@ impl ToolCompletionRequest {
self
}
/// Set stop sequences.
pub fn with_stop_sequences(mut self, stop_sequences: Vec<String>) -> Self {
self.stop_sequences = Some(stop_sequences);
self
}
/// Set tool choice mode.
pub fn with_tool_choice(mut self, choice: impl Into<String>) -> Self {
self.tool_choice = Some(choice.into());
@@ -504,8 +512,6 @@ pub fn strip_unsupported_completion_params(
/// This is the single helper function used by all providers to remove
/// parameters they don't support from tool calls, replacing duplicate stringly-typed logic.
///
/// Note: Only `Temperature` and `MaxTokens` are supported in `ToolCompletionRequest`.
/// `StopSequences` is only available in `CompletionRequest` and is not applicable to tool calls.
pub fn strip_unsupported_tool_params(
unsupported: &std::collections::HashSet<String>,
req: &mut ToolCompletionRequest,
@@ -519,7 +525,9 @@ pub fn strip_unsupported_tool_params(
if unsupported.contains(UnsupportedParam::MaxTokens.name()) {
req.max_tokens = None;
}
// Note: StopSequences is not a field in ToolCompletionRequest, so no action needed
if unsupported.contains(UnsupportedParam::StopSequences.name()) {
req.stop_sequences = None;
}
}
#[cfg(test)]
@@ -651,4 +659,17 @@ mod tests {
assert!(messages[2].tool_call_id.is_none());
assert!(messages[2].name.is_none());
}
#[test]
fn test_strip_unsupported_tool_params_strips_stop_sequences() {
let mut unsupported = std::collections::HashSet::new();
unsupported.insert(UnsupportedParam::StopSequences.name().to_string());
let mut req = ToolCompletionRequest::new(vec![ChatMessage::user("hello")], vec![]);
req.stop_sequences = Some(vec!["STOP".to_string()]);
strip_unsupported_tool_params(&unsupported, &mut req);
assert!(req.stop_sequences.is_none()); // safety: test assertion for explicit strip behavior
}
}
+1
View File
@@ -548,6 +548,7 @@ mod tests {
model: None,
max_tokens: None,
temperature: None,
stop_sequences: None,
tool_choice: None,
metadata: Default::default(),
};
+10 -1
View File
@@ -920,7 +920,16 @@ async fn async_main() -> anyhow::Result<()> {
}
if let Some(ref ws_arc) = webhook_server {
ws_arc.lock().await.shutdown().await;
let (shutdown_tx, handle) = {
let mut ws = ws_arc.lock().await;
ws.begin_shutdown()
};
if let Some(tx) = shutdown_tx {
let _ = tx.send(());
}
if let Some(handle) = handle {
let _ = handle.await;
}
}
if let Some(tunnel) = active_tunnel {
+1
View File
@@ -176,6 +176,7 @@ async fn llm_complete_with_tools(
model: req.model,
max_tokens: req.max_tokens,
temperature: req.temperature,
stop_sequences: req.stop_sequences,
tool_choice: req.tool_choice,
metadata: std::collections::HashMap::new(),
};
+2
View File
@@ -65,6 +65,7 @@ pub struct ProxyToolCompletionRequest {
pub model: Option<String>,
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
pub stop_sequences: Option<Vec<String>>,
pub tool_choice: Option<String>,
}
@@ -251,6 +252,7 @@ impl WorkerHttpClient {
model: request.model.clone(),
max_tokens: request.max_tokens,
temperature: request.temperature,
stop_sequences: request.stop_sequences.clone(),
tool_choice: request.tool_choice.clone(),
};
+509
View File
@@ -0,0 +1,509 @@
//! Tests for batch loading routine concurrent counts (N+1 query fix).
//!
//! Verifies:
//! 1. Batch query returns correct counts for multiple routines
//! 2. Concurrent limit enforcement uses batch counts correctly
#[cfg(feature = "libsql")]
mod tests {
use std::sync::Arc;
use chrono::Utc;
use uuid::Uuid;
use ironclaw::agent::routine::{
Routine, RoutineAction, RoutineGuardrails, RoutineRun, RunStatus, Trigger,
};
use ironclaw::db::Database;
async fn create_test_db() -> (Arc<dyn Database>, tempfile::TempDir) {
use ironclaw::db::libsql::LibSqlBackend;
let temp_dir = tempfile::tempdir().expect("tempdir");
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");
let db: Arc<dyn Database> = Arc::new(backend);
(db, temp_dir)
}
// -----------------------------------------------------------------------
// Test 1: Batch query returns correct counts for multiple routines
// -----------------------------------------------------------------------
#[tokio::test]
async fn batch_query_empty_list() {
let (db, _tmp) = create_test_db().await;
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");
}
#[tokio::test]
async fn batch_query_single_routine() {
let (db, _tmp) = create_test_db().await;
let routine_id = Uuid::new_v4();
// Create routine
let routine = Routine {
id: routine_id,
name: "test-routine".to_string(),
description: "Test".to_string(),
user_id: "default".to_string(),
enabled: true,
trigger: Trigger::Cron {
schedule: "* * * * *".to_string(),
timezone: None,
},
action: RoutineAction::Lightweight {
prompt: "test".to_string(),
context_paths: vec![],
max_tokens: 1000,
use_tools: false,
max_tool_rounds: 3,
},
guardrails: RoutineGuardrails {
cooldown: std::time::Duration::from_secs(0),
max_concurrent: 5,
dedup_window: None,
},
notify: Default::default(),
last_run_at: None,
next_fire_at: None,
run_count: 0,
consecutive_failures: 0,
state: serde_json::json!({}),
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
// Create 3 running runs
for _ in 0..3 {
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id,
trigger_type: "cron".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: None,
status: RunStatus::Running,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
}
// Batch query for single routine
let counts = db
.count_running_routine_runs_batch(&[routine_id])
.await
.expect("batch query should work");
assert_eq!(counts.len(), 1, "Should return 1 routine");
assert_eq!(counts[&routine_id], 3, "Should count 3 running runs");
}
#[tokio::test]
async fn batch_query_multiple_routines_different_counts() {
let (db, _tmp) = create_test_db().await;
let r1 = Uuid::new_v4();
let r2 = Uuid::new_v4();
let r3 = Uuid::new_v4();
// Create 3 routines
for routine_id in [r1, r2, r3] {
let routine = Routine {
id: routine_id,
name: format!("routine-{}", routine_id),
description: "Test".to_string(),
user_id: "default".to_string(),
enabled: true,
trigger: Trigger::Cron {
schedule: "* * * * *".to_string(),
timezone: None,
},
action: RoutineAction::Lightweight {
prompt: "test".to_string(),
context_paths: vec![],
max_tokens: 1000,
use_tools: false,
max_tool_rounds: 3,
},
guardrails: RoutineGuardrails {
cooldown: std::time::Duration::from_secs(0),
max_concurrent: 5,
dedup_window: None,
},
notify: Default::default(),
last_run_at: None,
next_fire_at: None,
run_count: 0,
consecutive_failures: 0,
state: serde_json::json!({}),
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
}
// r1: 2 running
for _ in 0..2 {
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id: r1,
trigger_type: "cron".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: None,
status: RunStatus::Running,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
}
// r2: 1 running
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id: r2,
trigger_type: "cron".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: None,
status: RunStatus::Running,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
// r3: 0 running (but has 1 Ok result)
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id: r3,
trigger_type: "cron".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: Some(Utc::now()),
status: RunStatus::Ok,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
// Single batch query for all 3
let counts = db
.count_running_routine_runs_batch(&[r1, r2, r3])
.await
.expect("batch query should work");
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[&r3], 0,
"r3 should have 0 running (Ok status is not running)"
);
}
#[tokio::test]
async fn batch_query_missing_routines_default_to_zero() {
let (db, _tmp) = create_test_db().await;
let r1 = Uuid::new_v4();
let r2 = Uuid::new_v4();
let r3 = Uuid::new_v4(); // This one won't exist
// Only create r1
let routine = Routine {
id: r1,
name: "routine-1".to_string(),
description: "Test".to_string(),
user_id: "default".to_string(),
enabled: true,
trigger: Trigger::Cron {
schedule: "* * * * *".to_string(),
timezone: None,
},
action: RoutineAction::Lightweight {
prompt: "test".to_string(),
context_paths: vec![],
max_tokens: 1000,
use_tools: false,
max_tool_rounds: 3,
},
guardrails: RoutineGuardrails {
cooldown: std::time::Duration::from_secs(0),
max_concurrent: 5,
dedup_window: None,
},
notify: Default::default(),
last_run_at: None,
next_fire_at: None,
run_count: 0,
consecutive_failures: 0,
state: serde_json::json!({}),
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
// r1 has 1 running
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id: r1,
trigger_type: "cron".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: None,
status: RunStatus::Running,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
// 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");
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");
}
#[tokio::test]
async fn batch_query_only_counts_running_status() {
let (db, _tmp) = create_test_db().await;
let routine_id = Uuid::new_v4();
// Create routine
let routine = Routine {
id: routine_id,
name: "test-routine".to_string(),
description: "Test".to_string(),
user_id: "default".to_string(),
enabled: true,
trigger: Trigger::Cron {
schedule: "* * * * *".to_string(),
timezone: None,
},
action: RoutineAction::Lightweight {
prompt: "test".to_string(),
context_paths: vec![],
max_tokens: 1000,
use_tools: false,
max_tool_rounds: 3,
},
guardrails: RoutineGuardrails {
cooldown: std::time::Duration::from_secs(0),
max_concurrent: 5,
dedup_window: None,
},
notify: Default::default(),
last_run_at: None,
next_fire_at: None,
run_count: 0,
consecutive_failures: 0,
state: serde_json::json!({}),
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
// Create 5 runs with mixed statuses
let statuses = [
RunStatus::Running,
RunStatus::Running,
RunStatus::Ok,
RunStatus::Failed,
RunStatus::Attention,
];
for status in statuses.iter() {
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id,
trigger_type: "cron".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: Some(Utc::now()),
status: *status,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
}
// Batch query should only count Running status
let counts = db
.count_running_routine_runs_batch(&[routine_id])
.await
.expect("batch query should work");
assert_eq!(
counts[&routine_id], 2,
"Should only count 2 Running status runs"
);
}
// -----------------------------------------------------------------------
// Test 2: Concurrent limit enforcement uses batch counts
// -----------------------------------------------------------------------
#[tokio::test]
async fn concurrent_limit_enforcement_with_batch_counts() {
let (db, _tmp) = create_test_db().await;
let r1 = Uuid::new_v4();
let r2 = Uuid::new_v4();
// Create 2 routines with max_concurrent=1 (r1) and max_concurrent=2 (r2)
for (routine_id, max_concurrent) in [(r1, 1), (r2, 2)] {
let routine = Routine {
id: routine_id,
name: format!("routine-{}", routine_id),
description: "Test".to_string(),
user_id: "default".to_string(),
enabled: true,
trigger: Trigger::Cron {
schedule: "* * * * *".to_string(),
timezone: None,
},
action: RoutineAction::Lightweight {
prompt: "test".to_string(),
context_paths: vec![],
max_tokens: 1000,
use_tools: false,
max_tool_rounds: 3,
},
guardrails: RoutineGuardrails {
cooldown: std::time::Duration::from_secs(0),
max_concurrent,
dedup_window: None,
},
notify: Default::default(),
last_run_at: None,
next_fire_at: None,
run_count: 0,
consecutive_failures: 0,
state: serde_json::json!({}),
created_at: Utc::now(),
updated_at: Utc::now(),
};
db.create_routine(&routine).await.expect("create routine");
}
// r1: create 1 running run (will hit max_concurrent=1)
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id: r1,
trigger_type: "cron".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: None,
status: RunStatus::Running,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
// r2: create 2 running runs (will hit max_concurrent=2)
for _ in 0..2 {
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id: r2,
trigger_type: "cron".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: None,
status: RunStatus::Running,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
}
// Batch query should return correct counts
let counts = db
.count_running_routine_runs_batch(&[r1, r2])
.await
.expect("batch query should work");
// Verify counts match the limits
assert_eq!(
counts[&r1], 1,
"r1 should have 1 running (at max_concurrent=1)"
);
assert_eq!(
counts[&r2], 2,
"r2 should have 2 running (at max_concurrent=2)"
);
// Now verify the limit enforcement logic
let r1_routine = db
.get_routine(r1)
.await
.expect("get routine")
.expect("routine exists");
let r2_routine = db
.get_routine(r2)
.await
.expect("get routine")
.expect("routine exists");
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");
// If we add one more run to r2, it should exceed limit
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id: r2,
trigger_type: "cron".to_string(),
trigger_detail: None,
started_at: Utc::now(),
completed_at: None,
status: RunStatus::Running,
result_summary: None,
tokens_used: None,
job_id: None,
created_at: Utc::now(),
};
db.create_routine_run(&run).await.expect("create run");
// Re-query to get updated counts
let counts = db
.count_running_routine_runs_batch(&[r1, r2])
.await
.expect("batch query should work");
let r2_exceeded_limit = counts[&r2] > r2_routine.guardrails.max_concurrent as i64;
assert!(r2_exceeded_limit, "r2 should have exceeded its limit");
}
}
@@ -0,0 +1,534 @@
"""
E2E tests for event-triggered routines with batch loading.
These tests verify that the N+1 query fix correctly:
1. Fires event-triggered routines on matching messages
2. Enforces concurrent limits via batch-loaded counts
3. Maintains performance with multiple simultaneous triggers
4. Works correctly through the full UI and agent loop
Playwright-based UI tests + SSE verification.
"""
import asyncio
import json
import pytest
from datetime import datetime, timedelta
from typing import List, Dict, Any
from playwright.async_api import async_playwright, Page, Browser, BrowserContext
@pytest.fixture
async def browser_and_context():
"""Create a Playwright browser and context for testing."""
async with async_playwright() as p:
browser = await p.chromium.launch(headless=True)
context = await browser.new_context()
yield browser, context
await context.close()
await browser.close()
class EventTriggerHelper:
"""Helper methods for event trigger testing."""
def __init__(self, page: Page):
self.page = page
async def navigate_to_routines(self):
"""Navigate to the routines page."""
await self.page.goto("http://localhost:8000/routines")
await self.page.wait_for_load_state("networkidle")
async def create_event_routine(
self,
name: str,
trigger_regex: str,
channel: str = "slack",
max_concurrent: int = 1,
) -> str:
"""
Create an event-triggered routine via UI.
Returns the routine ID.
"""
await self.navigate_to_routines()
# Click "New Routine" button
await self.page.click('button:has-text("New Routine")')
await self.page.wait_for_selector('input[name="routine_name"]')
# Fill routine details
await self.page.fill('input[name="routine_name"]', name)
await self.page.fill(
'textarea[name="routine_description"]',
f"Test routine: {name}",
)
# Select "Event Trigger" type
await self.page.click('label:has-text("Event Trigger")')
await self.page.wait_for_selector('input[name="trigger_regex"]')
# Fill trigger details
await self.page.fill('input[name="trigger_regex"]', trigger_regex)
await self.page.select_option('select[name="trigger_channel"]', channel)
# Set guardrails
await self.page.fill('input[name="max_concurrent"]', str(max_concurrent))
# Select lightweight action
await self.page.click('label:has-text("Lightweight")')
await self.page.fill(
'textarea[name="lightweight_prompt"]',
"Acknowledge the message and confirm trigger worked.",
)
# Save routine
await self.page.click('button:has-text("Save Routine")')
await self.page.wait_for_selector('text=Routine created successfully')
# Extract routine ID from success message or URL
routine_id = await self.page.locator('data-testid=routine-id').text_content()
return routine_id.strip() if routine_id else None
async def create_multiple_routines(
self, base_name: str, count: int, trigger_regex: str = None
) -> List[str]:
"""Create multiple event-triggered routines."""
routine_ids = []
for i in range(count):
name = f"{base_name}_{i}"
regex = trigger_regex or f"({i}|{base_name})"
routine_id = await self.create_event_routine(name, regex)
routine_ids.append(routine_id)
await asyncio.sleep(0.1) # Small delay between creations
return routine_ids
async def send_chat_message(self, message: str) -> List[str]:
"""
Send a chat message and return SSE events received.
Captures all routine firing events.
"""
await self.page.goto("http://localhost:8000/chat")
await self.page.wait_for_selector('input[placeholder*="message"]', timeout=5000)
# Collect SSE events
sse_events = []
async def capture_sse(response):
"""Intercept SSE events."""
if "event-stream" in response.headers.get("content-type", ""):
text = await response.text()
for line in text.split("\n"):
if line.startswith("data:"):
try:
event = json.loads(line[5:])
sse_events.append(event)
except json.JSONDecodeError:
pass
self.page.on("response", capture_sse)
# Send message
await self.page.fill('input[placeholder*="message"]', message)
await self.page.press('input[placeholder*="message"]', "Enter")
# Wait for response
await self.page.wait_for_selector('text=Message processed', timeout=10000)
await asyncio.sleep(0.5) # Allow time for SSE events
self.page.remove_listener("response", capture_sse)
return sse_events
async def get_routine_execution_log(self, routine_id: str) -> List[Dict]:
"""Get execution log entries for a routine."""
await self.page.goto(f"http://localhost:8000/routines/{routine_id}/executions")
await self.page.wait_for_load_state("networkidle")
# Extract log entries from table
rows = await self.page.locator("tbody tr").all()
executions = []
for row in rows:
cells = await row.locator("td").all()
if len(cells) >= 3:
execution = {
"timestamp": await cells[0].text_content(),
"status": await cells[1].text_content(),
"details": await cells[2].text_content(),
}
executions.append(execution)
return executions
async def check_database_queries_in_logs(
self, max_queries_expected: int = 1
) -> int:
"""Check debug logs for database query count."""
await self.page.goto("http://localhost:8000/debug/logs?filter=database")
await self.page.wait_for_load_state("networkidle")
# Count batch queries
log_lines = await self.page.locator("tr:has-text('batch')").all()
batch_count = len(log_lines)
# Count individual COUNT queries (should be 0 after fix)
count_queries = await self.page.locator("tr:has-text('COUNT')").all()
count_query_count = len(count_queries)
return batch_count, count_query_count
# =============================================================================
# Tests
# =============================================================================
@pytest.mark.asyncio
async def test_create_event_trigger_routine(browser_and_context):
"""Test creating an event-triggered routine via UI."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
routine_id = await helper.create_event_routine(
name="Test Trigger",
trigger_regex="test|demo",
channel="slack",
max_concurrent=1,
)
assert routine_id is not None, "Routine ID should be returned"
assert len(routine_id) > 0, "Routine ID should not be empty"
finally:
await page.close()
@pytest.mark.asyncio
async def test_event_trigger_fires_on_matching_message(browser_and_context):
"""Test that event-triggered routine fires when message matches."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
# Create routine
routine_id = await helper.create_event_routine(
name="Alert Handler",
trigger_regex="urgent|critical|alert",
channel="slack",
)
# Send matching message
sse_events = await helper.send_chat_message("URGENT: Server down!")
# Verify routine fired (look for event in SSE stream)
routine_fired = any(
event.get("type") == "routine_fired" and event.get("routine_id") == routine_id
for event in sse_events
)
assert routine_fired, "Routine should fire on matching message"
# Check execution log
executions = await helper.get_routine_execution_log(routine_id)
assert len(executions) > 0, "Execution should be logged"
assert "success" in executions[0]["status"].lower()
finally:
await page.close()
@pytest.mark.asyncio
async def test_event_trigger_skips_non_matching_message(browser_and_context):
"""Test that event-triggered routine skips when message doesn't match."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
# Create routine
routine_id = await helper.create_event_routine(
name="Alert Handler",
trigger_regex="urgent|critical|alert",
channel="slack",
)
# Send non-matching message
sse_events = await helper.send_chat_message("Hello, how are you?")
# Verify routine did NOT fire
routine_fired = any(
event.get("type") == "routine_fired" and event.get("routine_id") == routine_id
for event in sse_events
)
assert not routine_fired, "Routine should not fire on non-matching message"
finally:
await page.close()
@pytest.mark.asyncio
async def test_multiple_routines_fire_on_matching_message(browser_and_context):
"""Test that multiple event-triggered routines fire on same message."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
# Create 3 overlapping routines
routine_ids = await helper.create_multiple_routines(
base_name="Handler", count=3, trigger_regex="alert|warning|error"
)
# Send matching message
sse_events = await helper.send_chat_message("ERROR: Database connection failed")
# Verify all 3 routines fired
fired_count = sum(
1
for event in sse_events
if event.get("type") == "routine_fired" and event.get("routine_id") in routine_ids
)
assert (
fired_count >= 3
), f"Expected all 3 routines to fire, got {fired_count}"
finally:
await page.close()
@pytest.mark.asyncio
async def test_concurrent_limit_prevents_additional_fires(browser_and_context):
"""Test that concurrent limit is enforced via batch counts."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
# Create routine with max_concurrent=1
routine_id = await helper.create_event_routine(
name="Limited Handler",
trigger_regex="process|task",
max_concurrent=1,
)
# Trigger first message
await helper.send_chat_message("Process message 1")
await asyncio.sleep(1)
# Check first execution logged
executions_1 = await helper.get_routine_execution_log(routine_id)
assert len(executions_1) >= 1
# Trigger second message while first is still running
sse_events = await helper.send_chat_message("Process message 2")
# Second routine should be skipped (concurrent limit)
routine_skipped = any(
event.get("type") == "routine_skipped"
and event.get("reason") == "max_concurrent_reached"
and event.get("routine_id") == routine_id
for event in sse_events
)
assert routine_skipped, "Routine should be skipped when concurrent limit reached"
finally:
await page.close()
@pytest.mark.asyncio
async def test_rapid_messages_with_multiple_triggers_efficiency(browser_and_context):
"""Test efficiency of batch loading with multiple rapid messages."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
# Create 5 overlapping routines
routine_ids = await helper.create_multiple_routines(
base_name="Rapid", count=5, trigger_regex="test|demo|check"
)
# Send 10 matching messages rapidly
for i in range(10):
message = f"test message {i}"
await helper.send_chat_message(message)
await asyncio.sleep(0.1)
# Check database logs for query efficiency
batch_count, count_query_count = await helper.check_database_queries_in_logs()
# After fix: should have ~10 batch queries (1 per message)
# Before fix: would have ~50 individual COUNT queries (5 routines × 10 messages)
assert (
count_query_count == 0
), f"Should have 0 individual COUNT queries after fix, got {count_query_count}"
assert (
batch_count <= 15
), f"Should have <=15 batch queries for 10 messages, got {batch_count}"
finally:
await page.close()
@pytest.mark.asyncio
async def test_channel_filter_applied_correctly(browser_and_context):
"""Test that channel filter prevents non-matching messages."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
# Create routine for Slack channel
slack_routine_id = await helper.create_event_routine(
name="Slack Handler",
trigger_regex="alert",
channel="slack",
)
# Simulate message from Telegram channel
# (Note: In real UI, would need to change channel context)
page.goto(
"http://localhost:8000/chat?channel=telegram"
) # Switch channel
await helper.send_chat_message("alert: something urgent")
# Routine should not fire (different channel)
executions = await helper.get_routine_execution_log(slack_routine_id)
# Check if any recent execution (last 5 min) exists
recent = [
e
for e in executions
if (datetime.now() - datetime.fromisoformat(e["timestamp"])).total_seconds()
< 300
]
assert (
len(recent) == 0
), "Routine should not fire for different channel"
finally:
await page.close()
@pytest.mark.asyncio
async def test_batch_query_failure_handling(browser_and_context):
"""Test graceful handling of batch query failures."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
# Create routine
routine_id = await helper.create_event_routine(
name="Error Handler",
trigger_regex="test",
)
# Simulate database error in logs (if possible with test hooks)
# For now, just verify error handling doesn't crash UI
await helper.send_chat_message("test message")
# Check that UI remains responsive
assert await page.locator("text=Message processed").is_visible()
finally:
await page.close()
@pytest.mark.asyncio
async def test_routine_execution_history_display(browser_and_context):
"""Test that execution history correctly displays routine firings."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
# Create routine
routine_id = await helper.create_event_routine(
name="History Test",
trigger_regex="test",
)
# Trigger routine 3 times
for i in range(3):
await helper.send_chat_message(f"test message {i}")
await asyncio.sleep(0.2)
# Check execution log
executions = await helper.get_routine_execution_log(routine_id)
assert len(executions) >= 3, "Should have at least 3 executions logged"
# Verify all are recent (within last 5 minutes)
for execution in executions[:3]:
timestamp = datetime.fromisoformat(execution["timestamp"])
age = datetime.now() - timestamp
assert age < timedelta(minutes=5), "Execution should be recent"
finally:
await page.close()
@pytest.mark.asyncio
async def test_concurrent_batch_loads_independent(browser_and_context):
"""Test that concurrent messages each get independent batch queries."""
browser, context = browser_and_context
page = await context.new_page()
helper = EventTriggerHelper(page)
try:
# Create 5 routines matching different patterns
r1_id = await helper.create_event_routine(
name="Pattern A", trigger_regex="alpha|alpha_only"
)
r2_id = await helper.create_event_routine(
name="Pattern B", trigger_regex="beta|beta_only"
)
r3_id = await helper.create_event_routine(
name="Pattern AB", trigger_regex="alpha|beta|common"
)
# Send overlapping messages
# Message 1: matches r1, r3
sse1 = await helper.send_chat_message("alpha common")
await asyncio.sleep(0.1)
# Message 2: matches r2, r3
sse2 = await helper.send_chat_message("beta common")
await asyncio.sleep(0.1)
# Verify correct routines fired
r1_fired_msg1 = any(
e.get("routine_id") == r1_id for e in sse1 if e.get("type") == "routine_fired"
)
r2_fired_msg2 = any(
e.get("routine_id") == r2_id for e in sse2 if e.get("type") == "routine_fired"
)
r3_fired_both = (
any(
e.get("routine_id") == r3_id for e in sse1 if e.get("type") == "routine_fired"
)
and any(
e.get("routine_id") == r3_id for e in sse2 if e.get("type") == "routine_fired"
)
)
assert r1_fired_msg1, "Routine 1 should fire on message 1"
assert r2_fired_msg2, "Routine 2 should fire on message 2"
assert r3_fired_both, "Routine 3 should fire on both messages"
finally:
await page.close()
# =============================================================================
# Integration with existing test patterns
# =============================================================================
if __name__ == "__main__":
# Run tests with: pytest tests/e2e/scenarios/test_routine_event_batch.py -v
pytest.main([__file__, "-v", "-s"])