mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
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:
Generated
+13
-13
@@ -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]]
|
||||
|
||||
@@ -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": [
|
||||
|
||||
@@ -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": [
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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()));
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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
|
||||
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>'
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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]
|
||||
|
||||
@@ -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
@@ -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);
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -548,6 +548,7 @@ mod tests {
|
||||
model: None,
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
stop_sequences: None,
|
||||
tool_choice: None,
|
||||
metadata: Default::default(),
|
||||
};
|
||||
|
||||
+10
-1
@@ -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 {
|
||||
|
||||
@@ -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(),
|
||||
};
|
||||
|
||||
@@ -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(),
|
||||
};
|
||||
|
||||
|
||||
@@ -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"])
|
||||
Reference in New Issue
Block a user