mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 23:50:17 +00:00
Merge remote-tracking branch 'origin/staging' into refactor/architectural-hardening
# Conflicts: # src/agent/routine.rs
This commit is contained in:
+46
-9
@@ -750,6 +750,20 @@ impl Agent {
|
||||
"Message details"
|
||||
);
|
||||
|
||||
// Internal messages (e.g. job-monitor notifications) are already
|
||||
// rendered text and should be forwarded directly to the user without
|
||||
// entering the normal user-input pipeline (LLM/tool loop).
|
||||
// The `is_internal` field and `into_internal()` setter are pub(crate),
|
||||
// so external channels cannot spoof this flag.
|
||||
if message.is_internal {
|
||||
tracing::debug!(
|
||||
message_id = %message.id,
|
||||
channel = %message.channel,
|
||||
"Forwarding internal message"
|
||||
);
|
||||
return Ok(Some(message.content.clone()));
|
||||
}
|
||||
|
||||
// Set message tool context for this turn (current channel and target)
|
||||
// For Signal, use signal_target from metadata (group:ID or phone number),
|
||||
// otherwise fall back to user_id
|
||||
@@ -838,19 +852,42 @@ impl Agent {
|
||||
};
|
||||
|
||||
if let Some(pending) = pending_auth {
|
||||
match &submission {
|
||||
Submission::UserInput { content } => {
|
||||
return self
|
||||
.process_auth_token(message, &pending, content, session, thread_id)
|
||||
.await;
|
||||
}
|
||||
_ => {
|
||||
// Any control submission (interrupt, undo, etc.) cancels auth mode
|
||||
if pending.is_expired() {
|
||||
// TTL exceeded — clear stale auth mode
|
||||
tracing::warn!(
|
||||
extension = %pending.extension_name,
|
||||
"Auth mode expired after TTL, clearing"
|
||||
);
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.pending_auth = None;
|
||||
}
|
||||
// Fall through to normal handling
|
||||
}
|
||||
// If this was a user message (possibly a pasted token), return an
|
||||
// explicit error instead of forwarding it to the LLM/history.
|
||||
if matches!(submission, Submission::UserInput { .. }) {
|
||||
return Ok(Some(format!(
|
||||
"Authentication for **{}** expired. Please try again.",
|
||||
pending.extension_name
|
||||
)));
|
||||
}
|
||||
// Control submissions (interrupt, undo, etc.) fall through to normal handling
|
||||
} else {
|
||||
match &submission {
|
||||
Submission::UserInput { content } => {
|
||||
return self
|
||||
.process_auth_token(message, &pending, content, session, thread_id)
|
||||
.await;
|
||||
}
|
||||
_ => {
|
||||
// Any control submission (interrupt, undo, etc.) cancels auth mode
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.pending_auth = None;
|
||||
}
|
||||
// Fall through to normal handling
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -143,6 +143,11 @@ impl Agent {
|
||||
JobContext::with_user(&message.user_id, "chat", "Interactive chat session");
|
||||
job_ctx.http_interceptor = self.deps.http_interceptor.clone();
|
||||
job_ctx.user_timezone = user_tz.name().to_string();
|
||||
job_ctx.metadata = serde_json::json!({
|
||||
"notify_channel": message.channel,
|
||||
"notify_user": message.user_id,
|
||||
"notify_thread_id": message.thread_id,
|
||||
});
|
||||
|
||||
// Build system prompts once for this turn. Two variants: with tools
|
||||
// (normal iterations) and without (force_text final iteration).
|
||||
|
||||
+65
-14
@@ -21,6 +21,14 @@ use uuid::Uuid;
|
||||
use crate::channels::IncomingMessage;
|
||||
use crate::events::DomainEvent as SseEvent;
|
||||
|
||||
/// Route context for forwarding job monitor events back to the user's channel.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobMonitorRoute {
|
||||
pub channel: String,
|
||||
pub user_id: String,
|
||||
pub thread_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Spawn a background task that watches for events from a specific job and
|
||||
/// injects assistant messages into the agent loop.
|
||||
///
|
||||
@@ -35,6 +43,7 @@ pub fn spawn_job_monitor(
|
||||
job_id: Uuid,
|
||||
mut event_rx: broadcast::Receiver<(Uuid, SseEvent)>,
|
||||
inject_tx: mpsc::Sender<IncomingMessage>,
|
||||
route: JobMonitorRoute,
|
||||
) -> JoinHandle<()> {
|
||||
let short_id = job_id.to_string()[..8].to_string();
|
||||
|
||||
@@ -50,11 +59,15 @@ pub fn spawn_job_monitor(
|
||||
|
||||
match event {
|
||||
SseEvent::JobMessage { role, content, .. } if role == "assistant" => {
|
||||
let msg = IncomingMessage::new(
|
||||
"job_monitor",
|
||||
"system",
|
||||
let mut msg = IncomingMessage::new(
|
||||
route.channel.clone(),
|
||||
route.user_id.clone(),
|
||||
format!("[Job {}] Claude Code: {}", short_id, content),
|
||||
);
|
||||
)
|
||||
.into_internal();
|
||||
if let Some(ref thread_id) = route.thread_id {
|
||||
msg = msg.with_thread(thread_id.clone());
|
||||
}
|
||||
if inject_tx.send(msg).await.is_err() {
|
||||
tracing::debug!(
|
||||
job_id = %short_id,
|
||||
@@ -64,14 +77,18 @@ pub fn spawn_job_monitor(
|
||||
}
|
||||
}
|
||||
SseEvent::JobResult { status, .. } => {
|
||||
let msg = IncomingMessage::new(
|
||||
"job_monitor",
|
||||
"system",
|
||||
let mut msg = IncomingMessage::new(
|
||||
route.channel.clone(),
|
||||
route.user_id.clone(),
|
||||
format!(
|
||||
"[Job {}] Container finished (status: {})",
|
||||
short_id, status
|
||||
),
|
||||
);
|
||||
)
|
||||
.into_internal();
|
||||
if let Some(ref thread_id) = route.thread_id {
|
||||
msg = msg.with_thread(thread_id.clone());
|
||||
}
|
||||
let _ = inject_tx.send(msg).await;
|
||||
tracing::debug!(
|
||||
job_id = %short_id,
|
||||
@@ -108,13 +125,21 @@ pub fn spawn_job_monitor(
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn test_route() -> JobMonitorRoute {
|
||||
JobMonitorRoute {
|
||||
channel: "cli".to_string(),
|
||||
user_id: "user-1".to_string(),
|
||||
thread_id: Some("thread-1".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_monitor_forwards_assistant_messages() {
|
||||
let (event_tx, _) = broadcast::channel::<(Uuid, SseEvent)>(16);
|
||||
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
||||
|
||||
let job_id = Uuid::new_v4();
|
||||
let _handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx);
|
||||
let _handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx, test_route());
|
||||
|
||||
// Send an assistant message
|
||||
event_tx
|
||||
@@ -133,9 +158,11 @@ mod tests {
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(msg.channel, "job_monitor");
|
||||
assert_eq!(msg.user_id, "system");
|
||||
assert_eq!(msg.channel, "cli");
|
||||
assert_eq!(msg.user_id, "user-1");
|
||||
assert_eq!(msg.thread_id, Some("thread-1".to_string()));
|
||||
assert!(msg.content.contains("I found a bug"));
|
||||
assert!(msg.is_internal, "monitor messages must be marked internal");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -145,7 +172,7 @@ mod tests {
|
||||
|
||||
let job_id = Uuid::new_v4();
|
||||
let other_job_id = Uuid::new_v4();
|
||||
let _handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx);
|
||||
let _handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx, test_route());
|
||||
|
||||
// Send a message for a different job
|
||||
event_tx
|
||||
@@ -174,7 +201,7 @@ mod tests {
|
||||
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
||||
|
||||
let job_id = Uuid::new_v4();
|
||||
let handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx);
|
||||
let handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx, test_route());
|
||||
|
||||
// Send a completion event
|
||||
event_tx
|
||||
@@ -208,7 +235,7 @@ mod tests {
|
||||
let (inject_tx, mut inject_rx) = mpsc::channel::<IncomingMessage>(16);
|
||||
|
||||
let job_id = Uuid::new_v4();
|
||||
let _handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx);
|
||||
let _handle = spawn_job_monitor(job_id, event_tx.subscribe(), inject_tx, test_route());
|
||||
|
||||
// Send tool use event (should be skipped)
|
||||
event_tx
|
||||
@@ -242,4 +269,28 @@ mod tests {
|
||||
"should have timed out, no message expected"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression test: external channels must not be able to spoof the
|
||||
/// `is_internal` flag via metadata keys. A message created through
|
||||
/// the normal `IncomingMessage::new` + `with_metadata` path must
|
||||
/// always have `is_internal == false`, regardless of metadata content.
|
||||
#[test]
|
||||
fn test_external_metadata_cannot_spoof_internal_flag() {
|
||||
let msg = IncomingMessage::new("wasm_channel", "attacker", "pwned").with_metadata(
|
||||
serde_json::json!({
|
||||
"__internal_job_monitor": true,
|
||||
"is_internal": true,
|
||||
}),
|
||||
);
|
||||
assert!(
|
||||
!msg.is_internal,
|
||||
"with_metadata must not set is_internal — only into_internal() can"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_into_internal_sets_flag() {
|
||||
let msg = IncomingMessage::new("monitor", "system", "test").into_internal();
|
||||
assert!(msg.is_internal);
|
||||
}
|
||||
}
|
||||
|
||||
+127
-8
@@ -32,7 +32,9 @@ use crate::llm::{
|
||||
ChatMessage, CompletionRequest, FinishReason, LlmProvider, ToolCall, ToolCompletionRequest,
|
||||
};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::{ApprovalContext, ApprovalRequirement, ToolError, ToolRegistry};
|
||||
use crate::tools::{
|
||||
ApprovalContext, ApprovalRequirement, ToolError, ToolRegistry, prepare_tool_params,
|
||||
};
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
enum EventMatcher {
|
||||
@@ -139,6 +141,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 +192,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 +226,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 +306,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 +985,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 +1062,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,
|
||||
@@ -1034,13 +1120,14 @@ async fn execute_routine_tool(
|
||||
.get(&tc.name)
|
||||
.await
|
||||
.ok_or_else(|| format!("Tool '{}' not found", tc.name))?;
|
||||
let normalized_params = prepare_tool_params(tool.as_ref(), &tc.arguments);
|
||||
|
||||
// Check approval requirement: only allow Never tools in lightweight routines.
|
||||
// UnlessAutoApproved and Always tools are blocked to prevent prompt injection attacks.
|
||||
// Lightweight routines can be triggered by external events and may process untrusted data,
|
||||
// making them vulnerable to prompt injection that could trick the LLM into calling
|
||||
// sensitive tools. Blocking these tools entirely is the safest approach.
|
||||
match tool.requires_approval(&tc.arguments) {
|
||||
match tool.requires_approval(&normalized_params) {
|
||||
ApprovalRequirement::Never => {}
|
||||
ApprovalRequirement::UnlessAutoApproved | ApprovalRequirement::Always => {
|
||||
return Err(format!(
|
||||
@@ -1052,7 +1139,10 @@ async fn execute_routine_tool(
|
||||
}
|
||||
|
||||
// Validate tool parameters
|
||||
let validation = ctx.safety.validator().validate_tool_params(&tc.arguments);
|
||||
let validation = ctx
|
||||
.safety
|
||||
.validator()
|
||||
.validate_tool_params(&normalized_params);
|
||||
if !validation.is_valid {
|
||||
let details = validation
|
||||
.errors
|
||||
@@ -1067,7 +1157,7 @@ async fn execute_routine_tool(
|
||||
let timeout = tool.execution_timeout();
|
||||
let start = std::time::Instant::now();
|
||||
let result = tokio::time::timeout(timeout, async {
|
||||
tool.execute(tc.arguments.clone(), job_ctx).await
|
||||
tool.execute(normalized_params.clone(), job_ctx).await
|
||||
})
|
||||
.await;
|
||||
let elapsed = start.elapsed();
|
||||
@@ -1386,4 +1476,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
|
||||
}
|
||||
}
|
||||
|
||||
+117
-12
@@ -17,7 +17,7 @@ use crate::events::DomainEvent as SseEvent;
|
||||
use crate::hooks::HookRegistry;
|
||||
use crate::llm::LlmProvider;
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::{ApprovalContext, ToolRegistry};
|
||||
use crate::tools::{ApprovalContext, ToolRegistry, prepare_tool_params};
|
||||
use crate::worker::job::{Worker, WorkerDeps};
|
||||
|
||||
/// Message to send to a worker.
|
||||
@@ -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}"),
|
||||
@@ -505,8 +511,10 @@ impl Scheduler {
|
||||
.into());
|
||||
}
|
||||
|
||||
let normalized_params = prepare_tool_params(tool.as_ref(), ¶ms);
|
||||
|
||||
// Scheduler-specific approval check
|
||||
let requirement = tool.requires_approval(¶ms);
|
||||
let requirement = tool.requires_approval(&normalized_params);
|
||||
let blocked =
|
||||
ApprovalContext::is_blocked_or_default(&approval_context, tool_name, requirement);
|
||||
if blocked {
|
||||
@@ -518,7 +526,11 @@ impl Scheduler {
|
||||
|
||||
// Delegate to shared tool execution pipeline
|
||||
let output_str = crate::tools::execute::execute_tool_with_safety(
|
||||
&tools, &safety, tool_name, ¶ms, &job_ctx,
|
||||
&tools,
|
||||
&safety,
|
||||
tool_name,
|
||||
&normalized_params,
|
||||
&job_ctx,
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -832,6 +844,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
|
||||
@@ -1040,4 +1070,79 @@ mod tests {
|
||||
"hard_gate should pass with explicit permission"
|
||||
);
|
||||
}
|
||||
|
||||
struct NormalizedApprovalTool;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for NormalizedApprovalTool {
|
||||
fn name(&self) -> &str {
|
||||
"normalized_gate"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"approval depends on normalized params"
|
||||
}
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"safe": { "type": "boolean" }
|
||||
}
|
||||
})
|
||||
}
|
||||
async fn execute(
|
||||
&self,
|
||||
_params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
Ok(ToolOutput::text(
|
||||
"normalized_ok",
|
||||
std::time::Instant::now().elapsed(),
|
||||
))
|
||||
}
|
||||
fn requires_approval(&self, params: &serde_json::Value) -> ApprovalRequirement {
|
||||
if params.get("safe").and_then(|v| v.as_bool()) == Some(true) {
|
||||
ApprovalRequirement::Never
|
||||
} else {
|
||||
ApprovalRequirement::Always
|
||||
}
|
||||
}
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_execute_tool_task_normalizes_params_before_approval() {
|
||||
let registry = ToolRegistry::new();
|
||||
registry.register(Arc::new(NormalizedApprovalTool)).await;
|
||||
|
||||
let cm = Arc::new(ContextManager::new(5));
|
||||
let job_id = cm.create_job("test", "normalized approval").await.unwrap(); // safety: test-only setup
|
||||
cm.update_context(job_id, |ctx| ctx.transition_to(JobState::InProgress, None))
|
||||
.await
|
||||
.unwrap() // safety: test-only setup
|
||||
.unwrap(); // safety: test-only setup
|
||||
|
||||
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: false,
|
||||
}));
|
||||
|
||||
let result = Scheduler::execute_tool_task(
|
||||
Arc::new(registry),
|
||||
cm,
|
||||
safety,
|
||||
None,
|
||||
job_id,
|
||||
"normalized_gate",
|
||||
serde_json::json!({"safe": "true"}),
|
||||
)
|
||||
.await;
|
||||
|
||||
#[rustfmt::skip]
|
||||
assert!( // safety: test-only assertion
|
||||
result.is_ok(),
|
||||
"stringified boolean should normalize before approval: {result:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+50
-11
@@ -12,7 +12,7 @@
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use chrono::{DateTime, TimeDelta, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -92,8 +92,11 @@ impl Session {
|
||||
None => self.create_thread(),
|
||||
Some(id) => {
|
||||
if self.threads.contains_key(&id) {
|
||||
// Safe: contains_key confirmed the entry exists.
|
||||
self.threads.get_mut(&id).unwrap()
|
||||
// Entry existence confirmed by contains_key above.
|
||||
// get_mut borrows self.threads mutably, so we can't
|
||||
// combine the check and access into if-let without
|
||||
// conflicting with the self.create_thread() fallback.
|
||||
self.threads.get_mut(&id).unwrap() // safety: contains_key guard above
|
||||
} else {
|
||||
// Stale active_thread ID: create a new thread, which
|
||||
// updates self.active_thread to the new thread's ID.
|
||||
@@ -132,6 +135,12 @@ pub enum ThreadState {
|
||||
|
||||
/// Pending auth token request.
|
||||
///
|
||||
/// Auth mode TTL — must stay in sync with
|
||||
/// `crate::cli::oauth_defaults::OAUTH_FLOW_EXPIRY` (5 minutes / 300 s).
|
||||
/// Defined separately to avoid a session→cli module dependency.
|
||||
const AUTH_MODE_TTL_SECS: i64 = 300;
|
||||
const AUTH_MODE_TTL: TimeDelta = TimeDelta::seconds(AUTH_MODE_TTL_SECS);
|
||||
|
||||
/// When `tool_auth` returns `awaiting_token`, the thread enters auth mode.
|
||||
/// The next user message is intercepted before entering the normal pipeline
|
||||
/// (no logging, no turn creation, no history) and routed directly to the
|
||||
@@ -140,6 +149,16 @@ pub enum ThreadState {
|
||||
pub struct PendingAuth {
|
||||
/// Extension name to authenticate.
|
||||
pub extension_name: String,
|
||||
/// When this auth mode was entered. Used for TTL expiry.
|
||||
#[serde(default = "Utc::now")]
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl PendingAuth {
|
||||
/// Returns `true` if this auth mode has exceeded the TTL.
|
||||
pub fn is_expired(&self) -> bool {
|
||||
Utc::now() - self.created_at > AUTH_MODE_TTL
|
||||
}
|
||||
}
|
||||
|
||||
/// Pending tool approval request stored on a thread.
|
||||
@@ -295,7 +314,10 @@ impl Thread {
|
||||
/// Enter auth mode: next user message will be routed directly to
|
||||
/// the credential store, bypassing the normal pipeline entirely.
|
||||
pub fn enter_auth_mode(&mut self, extension_name: String) {
|
||||
self.pending_auth = Some(PendingAuth { extension_name });
|
||||
self.pending_auth = Some(PendingAuth {
|
||||
extension_name,
|
||||
created_at: Utc::now(),
|
||||
});
|
||||
self.updated_at = Utc::now();
|
||||
}
|
||||
|
||||
@@ -684,15 +706,16 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_enter_auth_mode() {
|
||||
let before = Utc::now();
|
||||
let mut thread = Thread::new(Uuid::new_v4());
|
||||
assert!(thread.pending_auth.is_none());
|
||||
|
||||
thread.enter_auth_mode("telegram".to_string());
|
||||
assert!(thread.pending_auth.is_some());
|
||||
assert_eq!(
|
||||
thread.pending_auth.as_ref().unwrap().extension_name,
|
||||
"telegram"
|
||||
);
|
||||
let pending = thread.pending_auth.as_ref().unwrap();
|
||||
assert_eq!(pending.extension_name, "telegram");
|
||||
assert!(pending.created_at >= before);
|
||||
assert!(!pending.is_expired());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -702,8 +725,9 @@ mod tests {
|
||||
|
||||
let pending = thread.take_pending_auth();
|
||||
assert!(pending.is_some());
|
||||
assert_eq!(pending.unwrap().extension_name, "notion");
|
||||
|
||||
let pending = pending.unwrap();
|
||||
assert_eq!(pending.extension_name, "notion");
|
||||
assert!(!pending.is_expired());
|
||||
// Should be cleared after take
|
||||
assert!(thread.pending_auth.is_none());
|
||||
assert!(thread.take_pending_auth().is_none());
|
||||
@@ -717,10 +741,25 @@ mod tests {
|
||||
let json = serde_json::to_string(&thread).expect("should serialize");
|
||||
assert!(json.contains("pending_auth"));
|
||||
assert!(json.contains("openai"));
|
||||
assert!(json.contains("created_at"));
|
||||
|
||||
let restored: Thread = serde_json::from_str(&json).expect("should deserialize");
|
||||
assert!(restored.pending_auth.is_some());
|
||||
assert_eq!(restored.pending_auth.unwrap().extension_name, "openai");
|
||||
let pending = restored.pending_auth.unwrap();
|
||||
assert_eq!(pending.extension_name, "openai");
|
||||
assert!(!pending.is_expired());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pending_auth_expiry() {
|
||||
let mut pending = PendingAuth {
|
||||
extension_name: "test".to_string(),
|
||||
created_at: Utc::now(),
|
||||
};
|
||||
assert!(!pending.is_expired());
|
||||
// Backdate beyond the TTL
|
||||
pending.created_at = Utc::now() - AUTH_MODE_TTL - TimeDelta::seconds(1);
|
||||
assert!(pending.is_expired());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
+24
-1
@@ -1540,7 +1540,8 @@ impl Agent {
|
||||
.configure_token(&pending.extension_name, token)
|
||||
.await
|
||||
{
|
||||
Ok(result) => {
|
||||
Ok(result) if result.activated => {
|
||||
// Ensure extension is actually activated
|
||||
tracing::info!(
|
||||
"Extension '{}' configured via auth mode: {}",
|
||||
pending.extension_name,
|
||||
@@ -1560,6 +1561,28 @@ impl Agent {
|
||||
.await;
|
||||
Ok(Some(result.message))
|
||||
}
|
||||
Ok(result) => {
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.enter_auth_mode(pending.extension_name.clone());
|
||||
}
|
||||
}
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::AuthRequired {
|
||||
extension_name: pending.extension_name.clone(),
|
||||
instructions: Some(result.message.clone()),
|
||||
auth_url: None,
|
||||
setup_url: None,
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
Ok(Some(result.message))
|
||||
}
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
// Token validation errors: re-enter auth mode and re-prompt
|
||||
|
||||
+1
-1
@@ -594,7 +594,7 @@ impl AppBuilder {
|
||||
let entries: Vec<_> = catalog
|
||||
.all()
|
||||
.iter()
|
||||
.map(|m| m.to_registry_entry())
|
||||
.filter_map(|m| m.to_registry_entry())
|
||||
.collect();
|
||||
tracing::debug!(
|
||||
count = entries.len(),
|
||||
|
||||
@@ -83,6 +83,11 @@ pub struct IncomingMessage {
|
||||
pub timezone: Option<String>,
|
||||
/// File or media attachments on this message.
|
||||
pub attachments: Vec<IncomingAttachment>,
|
||||
/// Internal-only flag: message was generated inside the process (e.g. job
|
||||
/// monitor) and must bypass the normal user-input pipeline. This field is
|
||||
/// **not** settable via `with_metadata()` — only trusted code paths inside
|
||||
/// the binary can set it, preventing external channels from spoofing it.
|
||||
pub(crate) is_internal: bool,
|
||||
}
|
||||
|
||||
impl IncomingMessage {
|
||||
@@ -103,6 +108,7 @@ impl IncomingMessage {
|
||||
metadata: serde_json::Value::Null,
|
||||
timezone: None,
|
||||
attachments: Vec::new(),
|
||||
is_internal: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -135,6 +141,12 @@ impl IncomingMessage {
|
||||
self.attachments = attachments;
|
||||
self
|
||||
}
|
||||
|
||||
/// Mark this message as internal (bypasses user-input pipeline).
|
||||
pub(crate) fn into_internal(mut self) -> Self {
|
||||
self.is_internal = true;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Stream of incoming messages.
|
||||
|
||||
+13
-13
@@ -140,7 +140,7 @@ struct WebhookRequest {
|
||||
content: String,
|
||||
/// Optional thread ID for conversation tracking.
|
||||
thread_id: Option<String>,
|
||||
/// Deprecated: webhook secret in request body. Use X-IronClaw-Signature header instead.
|
||||
/// Deprecated: webhook secret in request body. Use X-Hub-Signature-256 header instead.
|
||||
/// This field is accepted for backward compatibility but will be removed in a future release.
|
||||
secret: Option<String>,
|
||||
/// Whether to wait for a synchronous response.
|
||||
@@ -288,7 +288,7 @@ async fn webhook_handler(
|
||||
}
|
||||
};
|
||||
|
||||
match headers.get("x-ironclaw-signature") {
|
||||
match headers.get("x-hub-signature-256") {
|
||||
Some(raw_signature) => match raw_signature.to_str() {
|
||||
Ok(signature) => {
|
||||
if !verify_hmac_signature(expected_secret, &body, signature) {
|
||||
@@ -325,7 +325,7 @@ async fn webhook_handler(
|
||||
message_id: Uuid::nil(),
|
||||
status: "error".to_string(),
|
||||
response: Some(
|
||||
"Webhook authentication required. Provide X-IronClaw-Signature header \
|
||||
"Webhook authentication required. Provide X-Hub-Signature-256 header \
|
||||
(preferred) or 'secret' field in body (deprecated)."
|
||||
.to_string(),
|
||||
),
|
||||
@@ -341,7 +341,7 @@ async fn webhook_handler(
|
||||
{
|
||||
tracing::warn!(
|
||||
"Webhook authenticated via deprecated 'secret' field in request body. \
|
||||
Migrate to X-IronClaw-Signature header (HMAC-SHA256). \
|
||||
Migrate to X-Hub-Signature-256 header (HMAC-SHA256). \
|
||||
Body secret support will be removed in a future release."
|
||||
);
|
||||
fallback_req = Some(req);
|
||||
@@ -364,7 +364,7 @@ async fn webhook_handler(
|
||||
message_id: Uuid::nil(),
|
||||
status: "error".to_string(),
|
||||
response: Some(
|
||||
"Webhook authentication required. Provide X-IronClaw-Signature header \
|
||||
"Webhook authentication required. Provide X-Hub-Signature-256 header \
|
||||
(preferred) or 'secret' field in body (deprecated)."
|
||||
.to_string(),
|
||||
),
|
||||
@@ -726,7 +726,7 @@ mod tests {
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-ironclaw-signature", signature)
|
||||
.header("x-hub-signature-256", signature)
|
||||
.body(Body::from(body_bytes))
|
||||
.unwrap();
|
||||
|
||||
@@ -749,7 +749,7 @@ mod tests {
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-ironclaw-signature", signature)
|
||||
.header("x-hub-signature-256", signature)
|
||||
.body(Body::from(body_bytes))
|
||||
.unwrap();
|
||||
|
||||
@@ -770,7 +770,7 @@ mod tests {
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-ironclaw-signature", "not-a-valid-signature")
|
||||
.header("x-hub-signature-256", "not-a-valid-signature")
|
||||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||
.unwrap();
|
||||
|
||||
@@ -919,7 +919,7 @@ mod tests {
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-ironclaw-signature", signature)
|
||||
.header("x-hub-signature-256", signature)
|
||||
.body(Body::from(body_bytes))
|
||||
.unwrap();
|
||||
|
||||
@@ -941,7 +941,7 @@ mod tests {
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-ironclaw-signature", signature)
|
||||
.header("x-hub-signature-256", signature)
|
||||
.body(Body::from(body))
|
||||
.unwrap();
|
||||
|
||||
@@ -966,7 +966,7 @@ mod tests {
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "text/plain")
|
||||
.header("x-ironclaw-signature", signature)
|
||||
.header("x-hub-signature-256", signature)
|
||||
.body(Body::from(body_bytes))
|
||||
.unwrap();
|
||||
|
||||
@@ -991,7 +991,7 @@ mod tests {
|
||||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||
.unwrap();
|
||||
req.headers_mut().insert(
|
||||
"x-ironclaw-signature",
|
||||
"x-hub-signature-256",
|
||||
HeaderValue::from_bytes(b"\xFF").unwrap(),
|
||||
);
|
||||
|
||||
@@ -1083,7 +1083,7 @@ mod tests {
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-ironclaw-signature", signature)
|
||||
.header("x-hub-signature-256", signature)
|
||||
.body(Body::from(body_bytes))
|
||||
.unwrap();
|
||||
|
||||
|
||||
@@ -32,7 +32,7 @@ const MAX_HTTP_RESPONSE_SIZE: usize = 10 * 1024 * 1024;
|
||||
const MAX_REPLY_TARGETS: usize = 10000;
|
||||
const MAX_ERROR_LOG_BODY: usize = 1024;
|
||||
|
||||
const REPLY_TARGETS_CAP: NonZeroUsize = NonZeroUsize::new(MAX_REPLY_TARGETS).unwrap();
|
||||
const REPLY_TARGETS_CAP: NonZeroUsize = NonZeroUsize::new(MAX_REPLY_TARGETS).unwrap(); // safety: 10000 is nonzero
|
||||
|
||||
/// Recipient classification for outbound messages.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
|
||||
@@ -22,6 +22,7 @@ const KNOWN_CHANNELS: &[(&str, &str)] = &[
|
||||
("slack", "slack_channel"),
|
||||
("discord", "discord_channel"),
|
||||
("whatsapp", "whatsapp_channel"),
|
||||
("feishu", "feishu_channel"),
|
||||
];
|
||||
|
||||
/// Names of known channels that can be installed.
|
||||
|
||||
@@ -161,6 +161,13 @@ async fn register_channel(
|
||||
config_updates.insert("owner_id".to_string(), serde_json::json!(owner_id));
|
||||
}
|
||||
|
||||
// Inject channel-specific secrets into config for channels that need
|
||||
// credentials in API request bodies (e.g., Feishu token exchange).
|
||||
// The credential injection system only replaces placeholders in URLs
|
||||
// and headers, so channels like Feishu that exchange app_id + app_secret
|
||||
// for a tenant token need the raw values in their config.
|
||||
inject_channel_secrets_into_config(&channel_name, secrets_store, &mut config_updates).await;
|
||||
|
||||
if !config_updates.is_empty() {
|
||||
channel_arc.update_config(config_updates).await;
|
||||
tracing::info!(
|
||||
@@ -348,3 +355,62 @@ pub async fn inject_channel_credentials(
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Inject channel-specific secrets into the config JSON.
|
||||
///
|
||||
/// Some channels (e.g., Feishu) need raw credential values in their config
|
||||
/// because they perform token exchanges that require secrets in the HTTP
|
||||
/// request body. The standard credential injection system only replaces
|
||||
/// placeholders in URLs and headers, so this function fills config fields
|
||||
/// that map to secret names.
|
||||
///
|
||||
/// Mapping: for a channel named "feishu", secrets `feishu_app_id` and
|
||||
/// `feishu_app_secret` are injected as config keys `app_id` and `app_secret`.
|
||||
async fn inject_channel_secrets_into_config(
|
||||
channel_name: &str,
|
||||
secrets_store: &Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||
config_updates: &mut std::collections::HashMap<String, serde_json::Value>,
|
||||
) {
|
||||
// Map of (config_key, secret_name) pairs per channel.
|
||||
let secret_config_mappings: &[(&str, &str)] = match channel_name {
|
||||
"feishu" => &[
|
||||
("app_id", "feishu_app_id"),
|
||||
("app_secret", "feishu_app_secret"),
|
||||
],
|
||||
_ => return,
|
||||
};
|
||||
|
||||
let Some(secrets) = secrets_store else {
|
||||
return;
|
||||
};
|
||||
|
||||
for &(config_key, secret_name) in secret_config_mappings {
|
||||
match secrets.get_decrypted("default", secret_name).await {
|
||||
Ok(decrypted) => {
|
||||
config_updates.insert(
|
||||
config_key.to_string(),
|
||||
serde_json::Value::String(decrypted.expose().to_string()),
|
||||
);
|
||||
tracing::debug!(
|
||||
channel = %channel_name,
|
||||
config_key = %config_key,
|
||||
"Injected secret into channel config"
|
||||
);
|
||||
}
|
||||
Err(_) => {
|
||||
// Also try environment variable fallback.
|
||||
let env_name = secret_name.to_uppercase();
|
||||
if let Ok(val) = std::env::var(&env_name)
|
||||
&& !val.is_empty()
|
||||
{
|
||||
config_updates.insert(config_key.to_string(), serde_json::Value::String(val));
|
||||
tracing::debug!(
|
||||
channel = %channel_name,
|
||||
config_key = %config_key,
|
||||
"Injected secret from env into channel config"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()));
|
||||
|
||||
+111
-8
@@ -526,23 +526,33 @@ async fn oauth_callback_handler(
|
||||
.get("error_description")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| error.clone());
|
||||
clear_auth_mode(&state).await;
|
||||
return oauth_error_page(&description);
|
||||
}
|
||||
|
||||
let state_param = match params.get("state") {
|
||||
Some(s) if !s.is_empty() => s.clone(),
|
||||
_ => return oauth_error_page("IronClaw"),
|
||||
_ => {
|
||||
clear_auth_mode(&state).await;
|
||||
return oauth_error_page("IronClaw");
|
||||
}
|
||||
};
|
||||
|
||||
let code = match params.get("code") {
|
||||
Some(c) if !c.is_empty() => c.clone(),
|
||||
_ => return oauth_error_page("IronClaw"),
|
||||
_ => {
|
||||
clear_auth_mode(&state).await;
|
||||
return oauth_error_page("IronClaw");
|
||||
}
|
||||
};
|
||||
|
||||
// Look up the pending flow by CSRF state (atomic remove prevents replay)
|
||||
let ext_mgr = match state.extension_manager.as_ref() {
|
||||
Some(mgr) => mgr,
|
||||
None => return oauth_error_page("IronClaw"),
|
||||
None => {
|
||||
clear_auth_mode(&state).await;
|
||||
return oauth_error_page("IronClaw");
|
||||
}
|
||||
};
|
||||
|
||||
// Strip instance prefix from state for registry lookup.
|
||||
@@ -563,6 +573,7 @@ async fn oauth_callback_handler(
|
||||
lookup_key = %lookup_key,
|
||||
"OAuth callback received with unknown or expired state"
|
||||
);
|
||||
clear_auth_mode(&state).await;
|
||||
return oauth_error_page("IronClaw");
|
||||
}
|
||||
};
|
||||
@@ -581,6 +592,7 @@ async fn oauth_callback_handler(
|
||||
message: "OAuth flow expired. Please try again.".to_string(),
|
||||
});
|
||||
}
|
||||
clear_auth_mode(&state).await;
|
||||
return oauth_error_page(&flow.display_name);
|
||||
}
|
||||
|
||||
@@ -690,6 +702,10 @@ async fn oauth_callback_handler(
|
||||
}
|
||||
}
|
||||
|
||||
// Clear auth mode regardless of outcome so the next user message goes
|
||||
// through to the LLM instead of being intercepted as a token.
|
||||
clear_auth_mode(&state).await;
|
||||
|
||||
// After successful OAuth, auto-activate the extension so it moves
|
||||
// from "Installed (Authenticate)" → "Active" without a second click.
|
||||
// OAuth success is independent of activation — tokens are already stored.
|
||||
@@ -1147,7 +1163,7 @@ async fn chat_auth_token_handler(
|
||||
.configure_token(&req.extension_name, &req.token)
|
||||
.await
|
||||
{
|
||||
Ok(result) => {
|
||||
Ok(result) if result.activated => {
|
||||
// Clear auth mode on the active thread
|
||||
clear_auth_mode(&state).await;
|
||||
|
||||
@@ -1159,6 +1175,7 @@ async fn chat_auth_token_handler(
|
||||
|
||||
Ok(Json(ActionResponse::ok(result.message)))
|
||||
}
|
||||
Ok(result) => Ok(Json(ActionResponse::fail(result.message))),
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
// Re-emit auth_required for retry on validation errors
|
||||
@@ -2182,16 +2199,24 @@ async fn extensions_setup_submit_handler(
|
||||
"Extension manager not available (secrets store required)".to_string(),
|
||||
))?;
|
||||
|
||||
// Clear auth mode regardless of outcome so the next user message goes
|
||||
// through to the LLM instead of being intercepted as a token.
|
||||
clear_auth_mode(&state).await;
|
||||
|
||||
match ext_mgr.configure(&name, &req.secrets).await {
|
||||
Ok(result) => {
|
||||
// Broadcast auth_completed so the chat UI can dismiss any in-progress
|
||||
// auth card or setup modal that was triggered by tool_auth/tool_activate.
|
||||
// Broadcast completion status so chat UI can dismiss success cases while
|
||||
// leaving failed auth/configuration flows visible for correction.
|
||||
state.sse.broadcast(SseEvent::AuthCompleted {
|
||||
extension_name: name.clone(),
|
||||
success: true,
|
||||
success: result.activated,
|
||||
message: result.message.clone(),
|
||||
});
|
||||
let mut resp = ActionResponse::ok(result.message);
|
||||
let mut resp = if result.activated {
|
||||
ActionResponse::ok(result.message)
|
||||
} else {
|
||||
ActionResponse::fail(result.message)
|
||||
};
|
||||
resp.activated = Some(result.activated);
|
||||
resp.auth_url = result.auth_url;
|
||||
Ok(Json(resp))
|
||||
@@ -2346,12 +2371,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(),
|
||||
@@ -2832,6 +2861,80 @@ mod tests {
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_extensions_setup_submit_returns_failure_when_not_activated() {
|
||||
use axum::body::Body;
|
||||
use tower::ServiceExt;
|
||||
|
||||
let secrets = test_secrets_store();
|
||||
let (ext_mgr, _wasm_tools_dir, wasm_channels_dir) = test_ext_mgr(secrets);
|
||||
|
||||
let channel_name = "test-failing-channel";
|
||||
std::fs::write(
|
||||
wasm_channels_dir
|
||||
.path()
|
||||
.join(format!("{channel_name}.wasm")),
|
||||
b"\0asm fake",
|
||||
)
|
||||
.expect("write fake wasm");
|
||||
let caps = serde_json::json!({
|
||||
"type": "channel",
|
||||
"name": channel_name,
|
||||
"setup": {
|
||||
"required_secrets": [
|
||||
{"name": "BOT_TOKEN", "prompt": "Enter bot token"}
|
||||
]
|
||||
}
|
||||
});
|
||||
std::fs::write(
|
||||
wasm_channels_dir
|
||||
.path()
|
||||
.join(format!("{channel_name}.capabilities.json")),
|
||||
serde_json::to_string(&caps).expect("serialize caps"),
|
||||
)
|
||||
.expect("write capabilities");
|
||||
|
||||
let state = test_gateway_state(Some(ext_mgr));
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/api/extensions/{name}/setup",
|
||||
post(extensions_setup_submit_handler),
|
||||
)
|
||||
.with_state(state);
|
||||
|
||||
let req_body = serde_json::json!({
|
||||
"secrets": {
|
||||
"BOT_TOKEN": "dummy-token"
|
||||
}
|
||||
});
|
||||
let req = axum::http::Request::builder()
|
||||
.method("POST")
|
||||
.uri(format!("/api/extensions/{channel_name}/setup"))
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(req_body.to_string()))
|
||||
.expect("request");
|
||||
|
||||
let resp = ServiceExt::<axum::http::Request<Body>>::oneshot(app, req)
|
||||
.await
|
||||
.expect("response");
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
|
||||
let body = axum::body::to_bytes(resp.into_body(), 1024 * 64)
|
||||
.await
|
||||
.expect("body");
|
||||
let parsed: serde_json::Value = serde_json::from_slice(&body).expect("json response");
|
||||
assert_eq!(parsed["success"], serde_json::Value::Bool(false));
|
||||
assert_eq!(parsed["activated"], serde_json::Value::Bool(false));
|
||||
assert!(
|
||||
parsed["message"]
|
||||
.as_str()
|
||||
.unwrap_or_default()
|
||||
.contains("Activation failed"),
|
||||
"expected activation failure in message: {:?}",
|
||||
parsed
|
||||
);
|
||||
}
|
||||
|
||||
fn expired_flow_created_at() -> Option<std::time::Instant> {
|
||||
std::time::Instant::now()
|
||||
.checked_sub(oauth_defaults::OAUTH_FLOW_EXPIRY + std::time::Duration::from_secs(1))
|
||||
|
||||
@@ -19,6 +19,7 @@ let _loadThreadsTimer = null;
|
||||
const JOB_EVENTS_CAP = 500;
|
||||
const MEMORY_SEARCH_QUERY_MAX_LENGTH = 100;
|
||||
let stagedImages = [];
|
||||
let authFlowPending = false;
|
||||
let _ghostSuggestion = '';
|
||||
|
||||
// --- Slash Commands ---
|
||||
@@ -487,6 +488,12 @@ function clearSuggestionChips() {
|
||||
function sendMessage() {
|
||||
clearSuggestionChips();
|
||||
const input = document.getElementById('chat-input');
|
||||
if (authFlowPending) {
|
||||
showToast('Complete the auth step before sending chat messages.', 'info');
|
||||
const tokenField = document.querySelector('.auth-card .auth-token-input input');
|
||||
if (tokenField) tokenField.focus();
|
||||
return;
|
||||
}
|
||||
if (!currentThreadId) {
|
||||
console.warn('sendMessage: no thread selected, ignoring');
|
||||
return;
|
||||
@@ -515,7 +522,7 @@ function sendMessage() {
|
||||
}
|
||||
|
||||
function enableChatInput() {
|
||||
if (currentThreadIsReadOnly) return;
|
||||
if (currentThreadIsReadOnly || authFlowPending) return;
|
||||
const input = document.getElementById('chat-input');
|
||||
const btn = document.getElementById('send-btn');
|
||||
if (input) {
|
||||
@@ -600,6 +607,22 @@ document.getElementById('chat-input').addEventListener('paste', (e) => {
|
||||
}
|
||||
});
|
||||
|
||||
const chatMessagesEl = document.getElementById('chat-messages');
|
||||
chatMessagesEl.addEventListener('copy', (e) => {
|
||||
const selection = window.getSelection();
|
||||
if (!selection || selection.isCollapsed) return;
|
||||
const anchorNode = selection.anchorNode;
|
||||
const focusNode = selection.focusNode;
|
||||
if (!anchorNode || !focusNode) return;
|
||||
if (!chatMessagesEl.contains(anchorNode) || !chatMessagesEl.contains(focusNode)) return;
|
||||
const text = selection.toString();
|
||||
if (!text || !e.clipboardData) return;
|
||||
// Force plain-text clipboard output so dark-theme styling never leaks on paste.
|
||||
e.preventDefault();
|
||||
e.clipboardData.clearData();
|
||||
e.clipboardData.setData('text/plain', text);
|
||||
});
|
||||
|
||||
function addGeneratedImage(dataUrl, path) {
|
||||
const container = document.getElementById('chat-messages');
|
||||
const card = document.createElement('div');
|
||||
@@ -1182,6 +1205,7 @@ function showJobCard(data) {
|
||||
// --- Auth card ---
|
||||
|
||||
function handleAuthRequired(data) {
|
||||
setAuthFlowPending(true, data.instructions);
|
||||
if (data.auth_url) {
|
||||
// OAuth flow: show the global auth prompt with an OAuth button + optional token paste field.
|
||||
showAuthCard(data);
|
||||
@@ -1193,10 +1217,17 @@ function handleAuthRequired(data) {
|
||||
}
|
||||
|
||||
function handleAuthCompleted(data) {
|
||||
// Dismiss only the matching extension's UI so unrelated setup work is not interrupted.
|
||||
showToast(data.message, data.success ? 'success' : 'error');
|
||||
// Dismiss only the matching extension's UI so stale prompts are cleared.
|
||||
removeAuthCard(data.extension_name);
|
||||
closeConfigureModal(data.extension_name);
|
||||
showToast(data.message, data.success ? 'success' : 'error');
|
||||
if (!data.success) {
|
||||
setAuthFlowPending(false);
|
||||
if (currentTab === 'extensions') loadExtensions();
|
||||
enableChatInput();
|
||||
return;
|
||||
}
|
||||
setAuthFlowPending(false);
|
||||
if (shouldShowChannelConnectedMessage(data.extension_name, data.success)) {
|
||||
addMessage('system', 'Telegram is now connected. You can message me there and I can send you notifications.');
|
||||
}
|
||||
@@ -1376,6 +1407,7 @@ function cancelAuth(extensionName) {
|
||||
body: { extension_name: extensionName },
|
||||
}).catch(() => {});
|
||||
removeAuthCard(extensionName);
|
||||
setAuthFlowPending(false);
|
||||
enableChatInput();
|
||||
}
|
||||
|
||||
@@ -1393,6 +1425,24 @@ function showAuthCardError(extensionName, message) {
|
||||
}
|
||||
}
|
||||
|
||||
function setAuthFlowPending(pending, instructions) {
|
||||
authFlowPending = !!pending;
|
||||
const input = document.getElementById('chat-input');
|
||||
const btn = document.getElementById('send-btn');
|
||||
if (!input || !btn) return;
|
||||
if (authFlowPending) {
|
||||
input.disabled = true;
|
||||
btn.disabled = true;
|
||||
input.placeholder = instructions || 'Complete extension auth to continue chatting';
|
||||
return;
|
||||
}
|
||||
if (!currentThreadIsReadOnly) {
|
||||
input.disabled = false;
|
||||
btn.disabled = false;
|
||||
input.placeholder = I18n.t('chat.inputPlaceholder');
|
||||
}
|
||||
}
|
||||
|
||||
function loadHistory(before) {
|
||||
clearSuggestionChips();
|
||||
let historyUrl = '/api/chat/history?limit=50';
|
||||
@@ -1759,7 +1809,10 @@ chatInput.addEventListener('keydown', (e) => {
|
||||
}
|
||||
}
|
||||
|
||||
if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) {
|
||||
// Safari fires compositionend before keydown, so e.isComposing is already false
|
||||
// when Enter confirms IME input. keyCode 229 (VK_PROCESS) catches this case.
|
||||
// See https://bugs.webkit.org/show_bug.cgi?id=165004
|
||||
if (e.key === 'Enter' && !e.shiftKey && !e.isComposing && e.keyCode !== 229) {
|
||||
e.preventDefault();
|
||||
hideSlashAutocomplete();
|
||||
sendMessage();
|
||||
@@ -3535,10 +3588,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 +3662,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>'
|
||||
|
||||
@@ -595,6 +595,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>,
|
||||
@@ -607,25 +608,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 {
|
||||
@@ -647,6 +657,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()),
|
||||
@@ -678,6 +689,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;
|
||||
|
||||
+4
-1
@@ -405,7 +405,10 @@ fn check_routines_config() -> CheckResult {
|
||||
fn check_gateway_config(settings: &Settings) -> CheckResult {
|
||||
// Use the same resolve() path as runtime so invalid env values
|
||||
// (e.g. GATEWAY_PORT=abc) are caught here too.
|
||||
match crate::config::ChannelsConfig::resolve(settings) {
|
||||
let tunnel_enabled = crate::config::TunnelConfig::resolve(settings)
|
||||
.map(|t| t.is_enabled())
|
||||
.unwrap_or(false);
|
||||
match crate::config::ChannelsConfig::resolve(settings, tunnel_enabled) {
|
||||
Ok(channels) => match channels.gateway {
|
||||
Some(gw) => {
|
||||
if gw.auth_token.is_some() {
|
||||
|
||||
+587
@@ -0,0 +1,587 @@
|
||||
//! CLI command for viewing and managing gateway logs.
|
||||
//!
|
||||
//! Provides access to gateway logs through three mechanisms:
|
||||
//! - Reading the gateway log file (`~/.ironclaw/gateway.log`)
|
||||
//! - Streaming live logs via the gateway's SSE endpoint (`/api/logs/events`)
|
||||
//! - Getting/setting the runtime log level via `/api/logs/level`
|
||||
|
||||
use std::io::{Seek, SeekFrom};
|
||||
use std::path::Path;
|
||||
|
||||
use clap::Args;
|
||||
|
||||
/// View and manage gateway logs.
|
||||
#[derive(Args, Debug, Clone)]
|
||||
#[command(
|
||||
about = "View and manage gateway logs",
|
||||
long_about = "Tail gateway logs, stream live output, or adjust log level.\nExamples:\n ironclaw logs # Show last 200 lines\n ironclaw logs --follow # Stream live logs via SSE\n ironclaw logs --limit 50 --json # Last 50 lines as JSON\n ironclaw logs --level # Show current log level\n ironclaw logs --level debug # Set log level to debug"
|
||||
)]
|
||||
pub struct LogsCommand {
|
||||
/// Stream live logs from the running gateway via SSE.
|
||||
/// Replays recent history then streams new entries in real time.
|
||||
#[arg(short, long)]
|
||||
pub follow: bool,
|
||||
|
||||
/// Maximum number of lines to show (default: 200)
|
||||
#[arg(short, long, default_value = "200")]
|
||||
pub limit: usize,
|
||||
|
||||
/// Output log entries as JSON (one object per line)
|
||||
#[arg(long)]
|
||||
pub json: bool,
|
||||
|
||||
/// Display timestamps in local timezone
|
||||
#[arg(long)]
|
||||
pub local_time: bool,
|
||||
|
||||
/// Plain text output (no ANSI styling)
|
||||
#[arg(long)]
|
||||
pub plain: bool,
|
||||
|
||||
/// Gateway URL (default: http://{GATEWAY_HOST}:{GATEWAY_PORT})
|
||||
#[arg(long)]
|
||||
pub url: Option<String>,
|
||||
|
||||
/// Gateway auth token (reads GATEWAY_AUTH_TOKEN env if not set)
|
||||
#[arg(long)]
|
||||
pub token: Option<String>,
|
||||
|
||||
/// Connection timeout in milliseconds (default: 5000)
|
||||
#[arg(long, default_value = "5000")]
|
||||
pub timeout: u64,
|
||||
|
||||
/// Get or set runtime log level. Without a value, shows current level.
|
||||
/// With a value (trace|debug|info|warn|error), sets the level.
|
||||
#[arg(long, num_args = 0..=1, default_missing_value = "")]
|
||||
pub level: Option<String>,
|
||||
}
|
||||
|
||||
/// Resolved gateway connection parameters.
|
||||
struct GatewayParams {
|
||||
base_url: String,
|
||||
token: String,
|
||||
}
|
||||
|
||||
/// Run the logs CLI command.
|
||||
pub async fn run_logs_command(cmd: LogsCommand, config_path: Option<&Path>) -> anyhow::Result<()> {
|
||||
// --level takes priority: it's a control-plane operation, not log viewing.
|
||||
if let Some(level_arg) = &cmd.level {
|
||||
let params = resolve_gateway_params(&cmd, config_path).await?;
|
||||
if level_arg.is_empty() {
|
||||
return cmd_get_level(&cmd, ¶ms).await;
|
||||
} else {
|
||||
return cmd_set_level(&cmd, level_arg, ¶ms).await;
|
||||
}
|
||||
}
|
||||
|
||||
if cmd.follow {
|
||||
let params = resolve_gateway_params(&cmd, config_path).await?;
|
||||
cmd_follow(&cmd, ¶ms).await
|
||||
} else {
|
||||
cmd_show(&cmd)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Show log file ────────────────────────────────────────────────────────
|
||||
|
||||
/// Read the last N lines from `~/.ironclaw/gateway.log`.
|
||||
///
|
||||
/// Uses a reverse-scan strategy: seeks to the end of the file and reads
|
||||
/// backwards in chunks to find the last `limit` newlines, so memory usage
|
||||
/// is proportional to the output size, not the file size.
|
||||
fn cmd_show(cmd: &LogsCommand) -> anyhow::Result<()> {
|
||||
let log_path = crate::bootstrap::ironclaw_base_dir().join("gateway.log");
|
||||
if !log_path.exists() {
|
||||
anyhow::bail!(
|
||||
"No gateway log file found at {}.\n\
|
||||
The log file is created when the gateway runs in background mode \
|
||||
(e.g. `ironclaw gateway start`).",
|
||||
log_path.display()
|
||||
);
|
||||
}
|
||||
|
||||
let lines = tail_file(&log_path, cmd.limit)?;
|
||||
|
||||
if lines.is_empty() {
|
||||
println!("(log file is empty)");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if cmd.json {
|
||||
for line in &lines {
|
||||
let obj = serde_json::json!({ "line": line });
|
||||
println!("{}", obj);
|
||||
}
|
||||
} else {
|
||||
for line in &lines {
|
||||
println!("{}", line);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read the last `n` lines from a file by scanning backwards from EOF.
|
||||
///
|
||||
/// Reads in 8 KiB chunks from the end, counting newlines until enough
|
||||
/// are found or the beginning of the file is reached.
|
||||
fn tail_file(path: &Path, n: usize) -> anyhow::Result<Vec<String>> {
|
||||
let mut file = std::fs::File::open(path)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to open {}: {}", path.display(), e))?;
|
||||
|
||||
let file_len = file
|
||||
.seek(SeekFrom::End(0))
|
||||
.map_err(|e| anyhow::anyhow!("Failed to seek {}: {}", path.display(), e))?;
|
||||
|
||||
if file_len == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// Read backwards in chunks to find enough newlines.
|
||||
const CHUNK_SIZE: u64 = 8192;
|
||||
let mut tail_bytes = Vec::new();
|
||||
let mut newline_count = 0;
|
||||
let mut remaining = file_len;
|
||||
|
||||
while remaining > 0 && newline_count <= n {
|
||||
let read_size = std::cmp::min(CHUNK_SIZE, remaining);
|
||||
remaining -= read_size;
|
||||
|
||||
file.seek(SeekFrom::Start(remaining))
|
||||
.map_err(|e| anyhow::anyhow!("Seek failed: {e}"))?;
|
||||
|
||||
let mut chunk = vec![0u8; read_size as usize];
|
||||
std::io::Read::read_exact(&mut file, &mut chunk)
|
||||
.map_err(|e| anyhow::anyhow!("Read failed: {e}"))?;
|
||||
|
||||
// Count newlines in this chunk (backwards).
|
||||
for &byte in chunk.iter().rev() {
|
||||
if byte == b'\n' {
|
||||
newline_count += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Prepend chunk to collected bytes.
|
||||
chunk.append(&mut tail_bytes);
|
||||
tail_bytes = chunk;
|
||||
}
|
||||
|
||||
// Convert to string and take last N lines.
|
||||
let text = String::from_utf8_lossy(&tail_bytes);
|
||||
let all_lines: Vec<&str> = text.lines().collect();
|
||||
let start = all_lines.len().saturating_sub(n);
|
||||
|
||||
Ok(all_lines[start..].iter().map(|s| s.to_string()).collect())
|
||||
}
|
||||
|
||||
// ── Follow (live SSE stream) ─────────────────────────────────────────────
|
||||
|
||||
/// Connect to the gateway's `/api/logs/events` SSE endpoint and stream logs.
|
||||
async fn cmd_follow(cmd: &LogsCommand, params: &GatewayParams) -> anyhow::Result<()> {
|
||||
let timeout_dur = std::time::Duration::from_millis(cmd.timeout);
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.connect_timeout(timeout_dur)
|
||||
.build()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to create HTTP client: {e}"))?;
|
||||
|
||||
let url = format!("{}/api/logs/events", params.base_url);
|
||||
let resp = client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {}", params.token))
|
||||
.header("Accept", "text/event-stream")
|
||||
// No per-request timeout: SSE streams are long-lived.
|
||||
.timeout(std::time::Duration::from_secs(u64::MAX / 2))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to connect to gateway at {url}: {e}\n\
|
||||
Is the gateway running? Try `ironclaw gateway status`."
|
||||
)
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
anyhow::bail!(
|
||||
"Gateway returned HTTP {}: {}",
|
||||
resp.status(),
|
||||
resp.text().await.unwrap_or_default()
|
||||
);
|
||||
}
|
||||
|
||||
eprintln!("Connected to {} — streaming logs (Ctrl-C to stop)", url);
|
||||
|
||||
// Parse SSE stream line by line.
|
||||
let mut bytes_stream = resp.bytes_stream();
|
||||
let mut buffer = String::new();
|
||||
let mut lines_shown: usize = 0;
|
||||
|
||||
use futures::StreamExt;
|
||||
while let Some(chunk) = bytes_stream.next().await {
|
||||
let chunk = chunk.map_err(|e| anyhow::anyhow!("Stream error: {e}"))?;
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
// Process complete lines from the buffer.
|
||||
while let Some(newline_pos) = buffer.find('\n') {
|
||||
let line = buffer[..newline_pos].to_string(); // safety: find('\n') returns char boundary
|
||||
buffer = buffer[newline_pos + 1..].to_string(); // safety: '\n' is single byte
|
||||
|
||||
// SSE format: "data: {...}" lines carry the payload.
|
||||
if let Some(data) = line.strip_prefix("data: ")
|
||||
&& let Ok(entry) = serde_json::from_str::<serde_json::Value>(data)
|
||||
{
|
||||
print_log_entry(&entry, cmd);
|
||||
lines_shown += 1;
|
||||
}
|
||||
// Skip "event:", "id:", "retry:", and empty keepalive lines.
|
||||
}
|
||||
}
|
||||
|
||||
if lines_shown == 0 {
|
||||
eprintln!("(no log entries received)");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── Log level get/set ────────────────────────────────────────────────────
|
||||
|
||||
/// GET /api/logs/level — show the current log level.
|
||||
async fn cmd_get_level(cmd: &LogsCommand, params: &GatewayParams) -> anyhow::Result<()> {
|
||||
let timeout_dur = std::time::Duration::from_millis(cmd.timeout);
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(timeout_dur)
|
||||
.build()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to create HTTP client: {e}"))?;
|
||||
|
||||
let url = format!("{}/api/logs/level", params.base_url);
|
||||
let resp = client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {}", params.token))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to connect to gateway at {url}: {e}\n\
|
||||
Is the gateway running? Try `ironclaw gateway status`."
|
||||
)
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
anyhow::bail!(
|
||||
"Gateway returned HTTP {}: {}",
|
||||
resp.status(),
|
||||
resp.text().await.unwrap_or_default()
|
||||
);
|
||||
}
|
||||
|
||||
let body: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Invalid response: {e}"))?;
|
||||
|
||||
if cmd.json {
|
||||
println!(
|
||||
"{}",
|
||||
serde_json::to_string_pretty(&body).unwrap_or_default()
|
||||
);
|
||||
} else {
|
||||
let level = body
|
||||
.get("level")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown");
|
||||
println!("Current log level: {}", level);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// PUT /api/logs/level — change the runtime log level.
|
||||
async fn cmd_set_level(
|
||||
cmd: &LogsCommand,
|
||||
level: &str,
|
||||
params: &GatewayParams,
|
||||
) -> anyhow::Result<()> {
|
||||
const VALID: &[&str] = &["trace", "debug", "info", "warn", "error"];
|
||||
let level_lower = level.to_lowercase();
|
||||
if !VALID.contains(&level_lower.as_str()) {
|
||||
anyhow::bail!(
|
||||
"Invalid log level '{}'. Must be one of: {}",
|
||||
level,
|
||||
VALID.join(", ")
|
||||
);
|
||||
}
|
||||
|
||||
let timeout_dur = std::time::Duration::from_millis(cmd.timeout);
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(timeout_dur)
|
||||
.build()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to create HTTP client: {e}"))?;
|
||||
|
||||
let url = format!("{}/api/logs/level", params.base_url);
|
||||
let resp = client
|
||||
.put(&url)
|
||||
.header("Authorization", format!("Bearer {}", params.token))
|
||||
.json(&serde_json::json!({ "level": level_lower }))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to connect to gateway at {url}: {e}\n\
|
||||
Is the gateway running? Try `ironclaw gateway status`."
|
||||
)
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
anyhow::bail!(
|
||||
"Gateway returned HTTP {}: {}",
|
||||
resp.status(),
|
||||
resp.text().await.unwrap_or_default()
|
||||
);
|
||||
}
|
||||
|
||||
let body: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Invalid response: {e}"))?;
|
||||
|
||||
if cmd.json {
|
||||
println!(
|
||||
"{}",
|
||||
serde_json::to_string_pretty(&body).unwrap_or_default()
|
||||
);
|
||||
} else {
|
||||
let new_level = body
|
||||
.get("level")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(&level_lower);
|
||||
println!("Log level set to: {}", new_level);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ── Helpers ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// Resolve gateway connection params from CLI flags, config file, or env.
|
||||
///
|
||||
/// Priority: --url/--token flags > config TOML > env vars > defaults.
|
||||
async fn resolve_gateway_params(
|
||||
cmd: &LogsCommand,
|
||||
config_path: Option<&Path>,
|
||||
) -> anyhow::Result<GatewayParams> {
|
||||
// Load gateway config. Errors propagate when --config is explicit.
|
||||
let gw_config = load_gateway_config(config_path).await?;
|
||||
|
||||
// URL: --url flag > config TOML > env vars > defaults.
|
||||
let base_url = if let Some(url) = &cmd.url {
|
||||
url.trim_end_matches('/').to_string()
|
||||
} else if let Some(cfg) = &gw_config {
|
||||
format!("http://{}:{}", cfg.host, cfg.port)
|
||||
} else {
|
||||
let host = std::env::var("GATEWAY_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
|
||||
let port: u16 = std::env::var("GATEWAY_PORT")
|
||||
.ok()
|
||||
.and_then(|p| p.parse().ok())
|
||||
.unwrap_or(3000);
|
||||
format!("http://{}:{}", host, port)
|
||||
};
|
||||
|
||||
// Token: --token flag > config TOML > env var.
|
||||
let token = if let Some(token) = &cmd.token {
|
||||
token.clone()
|
||||
} else if let Some(t) = gw_config.as_ref().and_then(|c| c.auth_token.clone()) {
|
||||
t
|
||||
} else {
|
||||
std::env::var("GATEWAY_AUTH_TOKEN").map_err(|_| {
|
||||
anyhow::anyhow!(
|
||||
"No auth token provided. Use --token <TOKEN> or set GATEWAY_AUTH_TOKEN.\n\
|
||||
The token is printed when the gateway starts."
|
||||
)
|
||||
})?
|
||||
};
|
||||
|
||||
Ok(GatewayParams { base_url, token })
|
||||
}
|
||||
|
||||
/// Try to load gateway config from the TOML config file.
|
||||
///
|
||||
/// If `config_path` was explicitly provided (via `--config`), errors are
|
||||
/// propagated — the user asked for a specific file and deserves a clear
|
||||
/// failure when it is missing, unreadable, or malformed. When no path
|
||||
/// was given we fall back to env-only resolution and silently return
|
||||
/// `None` on failure so that `ironclaw logs` works without any config.
|
||||
async fn load_gateway_config(
|
||||
config_path: Option<&Path>,
|
||||
) -> anyhow::Result<Option<crate::config::GatewayConfig>> {
|
||||
if config_path.is_some() {
|
||||
// Explicit --config: propagate errors.
|
||||
let config = crate::config::Config::from_env_with_toml(config_path)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{e:#}"))?;
|
||||
Ok(config.channels.gateway)
|
||||
} else {
|
||||
// No explicit config: best-effort, swallow errors.
|
||||
let config = crate::config::Config::from_env_with_toml(None).await.ok();
|
||||
Ok(config.and_then(|c| c.channels.gateway))
|
||||
}
|
||||
}
|
||||
|
||||
/// Print a single log entry to stdout.
|
||||
fn print_log_entry(entry: &serde_json::Value, cmd: &LogsCommand) {
|
||||
if cmd.json {
|
||||
println!("{}", serde_json::to_string(entry).unwrap_or_default());
|
||||
return;
|
||||
}
|
||||
|
||||
let level = entry.get("level").and_then(|v| v.as_str()).unwrap_or("?");
|
||||
let target = entry.get("target").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let message = entry.get("message").and_then(|v| v.as_str()).unwrap_or("");
|
||||
let timestamp = entry
|
||||
.get("timestamp")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
let display_ts = if cmd.local_time {
|
||||
convert_to_local_time(timestamp)
|
||||
} else {
|
||||
timestamp.to_string()
|
||||
};
|
||||
|
||||
if cmd.plain {
|
||||
println!("{} {} [{}] {}", display_ts, level, target, message);
|
||||
} else {
|
||||
let level_colored = colorize_level(level);
|
||||
println!("{} {} [{}] {}", display_ts, level_colored, target, message);
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert an RFC 3339 timestamp to local time display.
|
||||
fn convert_to_local_time(ts: &str) -> String {
|
||||
chrono::DateTime::parse_from_rfc3339(ts)
|
||||
.map(|dt| {
|
||||
dt.with_timezone(&chrono::Local)
|
||||
.format("%Y-%m-%dT%H:%M:%S%.3f")
|
||||
.to_string()
|
||||
})
|
||||
.unwrap_or_else(|_| ts.to_string())
|
||||
}
|
||||
|
||||
/// Apply ANSI color to log level for terminal display.
|
||||
fn colorize_level(level: &str) -> String {
|
||||
match level {
|
||||
"ERROR" => format!("\x1b[31m{}\x1b[0m", level), // red
|
||||
"WARN" => format!("\x1b[33m{}\x1b[0m", level), // yellow
|
||||
"INFO" => format!("\x1b[32m{}\x1b[0m", level), // green
|
||||
"DEBUG" => format!("\x1b[36m{}\x1b[0m", level), // cyan
|
||||
"TRACE" => format!("\x1b[90m{}\x1b[0m", level), // gray
|
||||
_ => level.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_colorize_level() {
|
||||
assert!(colorize_level("ERROR").contains("\x1b[31m")); // safety: test-only
|
||||
assert!(colorize_level("WARN").contains("\x1b[33m")); // safety: test-only
|
||||
assert!(colorize_level("INFO").contains("\x1b[32m")); // safety: test-only
|
||||
assert!(colorize_level("DEBUG").contains("\x1b[36m")); // safety: test-only
|
||||
assert!(colorize_level("TRACE").contains("\x1b[90m")); // safety: test-only
|
||||
assert_eq!(colorize_level("UNKNOWN"), "UNKNOWN"); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_to_local_time_valid() {
|
||||
let ts = "2024-01-15T10:30:00.000Z";
|
||||
let result = convert_to_local_time(ts);
|
||||
assert!(result.contains("2024-01-15")); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_convert_to_local_time_invalid() {
|
||||
let ts = "not-a-timestamp";
|
||||
assert_eq!(convert_to_local_time(ts), "not-a-timestamp"); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_print_log_entry_json() {
|
||||
let entry = serde_json::json!({
|
||||
"level": "INFO",
|
||||
"target": "ironclaw::agent",
|
||||
"message": "test message",
|
||||
"timestamp": "2024-01-15T10:30:00.000Z"
|
||||
});
|
||||
let cmd = LogsCommand {
|
||||
follow: false,
|
||||
limit: 200,
|
||||
json: true,
|
||||
local_time: false,
|
||||
plain: false,
|
||||
url: None,
|
||||
token: None,
|
||||
timeout: 5000,
|
||||
level: None,
|
||||
};
|
||||
// Should not panic
|
||||
print_log_entry(&entry, &cmd);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tail_file_small() {
|
||||
let dir = tempfile::tempdir().unwrap(); // safety: test-only
|
||||
let path = dir.path().join("test.log");
|
||||
std::fs::write(&path, "line1\nline2\nline3\nline4\nline5\n").unwrap(); // safety: test-only
|
||||
|
||||
let result = tail_file(&path, 3).unwrap(); // safety: test-only
|
||||
assert_eq!(result, vec!["line3", "line4", "line5"]); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tail_file_fewer_lines_than_limit() {
|
||||
let dir = tempfile::tempdir().unwrap(); // safety: test-only
|
||||
let path = dir.path().join("test.log");
|
||||
std::fs::write(&path, "a\nb\n").unwrap(); // safety: test-only
|
||||
|
||||
let result = tail_file(&path, 200).unwrap(); // safety: test-only
|
||||
assert_eq!(result, vec!["a", "b"]); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tail_file_empty() {
|
||||
let dir = tempfile::tempdir().unwrap(); // safety: test-only
|
||||
let path = dir.path().join("test.log");
|
||||
std::fs::write(&path, "").unwrap(); // safety: test-only
|
||||
|
||||
let result = tail_file(&path, 10).unwrap(); // safety: test-only
|
||||
assert!(result.is_empty()); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tail_file_large() {
|
||||
let dir = tempfile::tempdir().unwrap(); // safety: test-only
|
||||
let path = dir.path().join("big.log");
|
||||
// Write 10000 lines to test chunked reading.
|
||||
let content: String = (0..10000).map(|i| format!("line {}\n", i)).collect();
|
||||
std::fs::write(&path, &content).unwrap(); // safety: test-only
|
||||
|
||||
let result = tail_file(&path, 5).unwrap(); // safety: test-only
|
||||
assert_eq!(result.len(), 5); // safety: test-only
|
||||
assert_eq!(result[0], "line 9995"); // safety: test-only
|
||||
assert_eq!(result[4], "line 9999"); // safety: test-only
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tail_file_no_trailing_newline() {
|
||||
let dir = tempfile::tempdir().unwrap(); // safety: test-only
|
||||
let path = dir.path().join("test.log");
|
||||
std::fs::write(&path, "line1\nline2\nline3").unwrap(); // safety: test-only
|
||||
|
||||
let result = tail_file(&path, 2).unwrap(); // safety: test-only
|
||||
assert_eq!(result, vec!["line2", "line3"]); // safety: test-only
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@
|
||||
//! - Managing OS service (`service install`, `service start`, `service stop`)
|
||||
//! - Listing configured channels (`channels list`)
|
||||
//! - Active health diagnostics (`doctor`)
|
||||
//! - Viewing gateway logs (`logs`)
|
||||
//! - Checking system health (`status`)
|
||||
|
||||
mod channels;
|
||||
@@ -19,6 +20,7 @@ mod config;
|
||||
mod doctor;
|
||||
#[cfg(feature = "import")]
|
||||
pub mod import;
|
||||
mod logs;
|
||||
mod mcp;
|
||||
pub mod memory;
|
||||
pub mod oauth_defaults;
|
||||
@@ -36,6 +38,7 @@ pub use config::{ConfigCommand, run_config_command};
|
||||
pub use doctor::run_doctor_command;
|
||||
#[cfg(feature = "import")]
|
||||
pub use import::{ImportCommand, run_import_command};
|
||||
pub use logs::{LogsCommand, run_logs_command};
|
||||
pub use mcp::{McpCommand, run_mcp_command};
|
||||
pub use memory::MemoryCommand;
|
||||
pub use memory::run_memory_command_with_db;
|
||||
@@ -206,6 +209,13 @@ pub enum Command {
|
||||
)]
|
||||
Doctor,
|
||||
|
||||
/// View and manage gateway logs
|
||||
#[command(
|
||||
about = "View and manage gateway logs",
|
||||
long_about = "Tail gateway logs, stream live output, or adjust log level.\nExamples:\n ironclaw logs # Show last 200 lines from gateway.log\n ironclaw logs --follow # Stream live logs via SSE\n ironclaw logs --level # Show current log level\n ironclaw logs --level debug # Set log level to debug"
|
||||
)]
|
||||
Logs(LogsCommand),
|
||||
|
||||
/// Show system health and diagnostics
|
||||
#[command(
|
||||
about = "Show system status",
|
||||
|
||||
+18
-6
@@ -127,7 +127,11 @@ fn cmd_list(
|
||||
.unwrap_or("none");
|
||||
println!(
|
||||
"{:<20} {:<8} {:<8} {:<10} {}",
|
||||
m.name, m.kind, m.version, auth, m.description
|
||||
m.name,
|
||||
m.kind,
|
||||
m.version.as_deref().unwrap_or("-"),
|
||||
auth,
|
||||
m.description
|
||||
);
|
||||
} else {
|
||||
println!("{:<20} {:<8} {}", m.name, m.kind, m.description);
|
||||
@@ -173,17 +177,25 @@ fn cmd_info(catalog: &RegistryCatalog, name: &str) -> anyhow::Result<()> {
|
||||
.map_err(|e| anyhow::anyhow!("{}", e))?;
|
||||
|
||||
println!("{} ({})", manifest.display_name, manifest.kind);
|
||||
println!(" Version: {}", manifest.version);
|
||||
if let Some(ref version) = manifest.version {
|
||||
println!(" Version: {}", version);
|
||||
}
|
||||
println!(" {}", manifest.description);
|
||||
|
||||
if !manifest.keywords.is_empty() {
|
||||
println!(" Keywords: {}", manifest.keywords.join(", "));
|
||||
}
|
||||
|
||||
println!("\nSource:");
|
||||
println!(" Directory: {}", manifest.source.dir);
|
||||
println!(" Crate: {}", manifest.source.crate_name);
|
||||
println!(" Capabilities: {}", manifest.source.capabilities);
|
||||
if let Some(ref source) = manifest.source {
|
||||
println!("\nSource:");
|
||||
println!(" Directory: {}", source.dir);
|
||||
println!(" Crate: {}", source.crate_name);
|
||||
println!(" Capabilities: {}", source.capabilities);
|
||||
}
|
||||
|
||||
if let Some(ref url) = manifest.url {
|
||||
println!("\nMCP Server URL: {}", url);
|
||||
}
|
||||
|
||||
if let Some(artifact) = manifest.artifacts.get("wasm32-wasip2") {
|
||||
println!("\nArtifact (wasm32-wasip2):");
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
---
|
||||
source: src/cli/mod.rs
|
||||
expression: help
|
||||
---
|
||||
Secure personal AI assistant that protects your data and expands its capabilities
|
||||
|
||||
Usage: ironclaw [OPTIONS] [COMMAND]
|
||||
|
||||
Commands:
|
||||
run Run the AI agent
|
||||
onboard Run interactive setup wizard
|
||||
config Manage app configs
|
||||
tool Manage WASM tools
|
||||
registry Browse/install extensions
|
||||
channels Manage channels
|
||||
routines Manage routines
|
||||
mcp Manage MCP servers
|
||||
memory Manage workspace memory
|
||||
pairing Manage DM pairing
|
||||
service Manage OS service
|
||||
skills Manage skills
|
||||
doctor Run diagnostics
|
||||
logs View and manage gateway logs
|
||||
status Show system status
|
||||
completion Generate completions
|
||||
import Import from other AI systems
|
||||
help Print this message or the help of the given subcommand(s)
|
||||
|
||||
Options:
|
||||
--cli-only Run in interactive CLI mode only (disable other channels)
|
||||
--no-db Skip database connection (for testing)
|
||||
-m, --message <MESSAGE> Single message mode - send one message and exit
|
||||
-c, --config <CONFIG> Configuration file path (optional, uses env vars by default)
|
||||
--no-onboard Skip first-run onboarding check
|
||||
-h, --help Print help (see more with '--help')
|
||||
-V, --version Print version
|
||||
@@ -20,6 +20,7 @@ Commands:
|
||||
service Manage OS service
|
||||
skills Manage skills
|
||||
doctor Run diagnostics
|
||||
logs View and manage gateway logs
|
||||
status Show system status
|
||||
completion Generate completions
|
||||
help Print this message or the help of the given subcommand(s)
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
---
|
||||
source: src/cli/mod.rs
|
||||
expression: help
|
||||
---
|
||||
IronClaw is a secure AI assistant. Use 'ironclaw <subcommand> --help' for details.
|
||||
Examples:
|
||||
ironclaw run # Start the agent
|
||||
ironclaw config list # List configs
|
||||
|
||||
Usage: ironclaw [OPTIONS] [COMMAND]
|
||||
|
||||
Commands:
|
||||
run Run the AI agent
|
||||
onboard Run interactive setup wizard
|
||||
config Manage app configs
|
||||
tool Manage WASM tools
|
||||
registry Browse/install extensions
|
||||
channels Manage channels
|
||||
routines Manage routines
|
||||
mcp Manage MCP servers
|
||||
memory Manage workspace memory
|
||||
pairing Manage DM pairing
|
||||
service Manage OS service
|
||||
skills Manage skills
|
||||
doctor Run diagnostics
|
||||
logs View and manage gateway logs
|
||||
status Show system status
|
||||
completion Generate completions
|
||||
import Import from other AI systems
|
||||
help Print this message or the help of the given subcommand(s)
|
||||
|
||||
Options:
|
||||
--cli-only
|
||||
Run in interactive CLI mode only (disable other channels)
|
||||
|
||||
--no-db
|
||||
Skip database connection (for testing)
|
||||
|
||||
-m, --message <MESSAGE>
|
||||
Single message mode - send one message and exit
|
||||
|
||||
-c, --config <CONFIG>
|
||||
Configuration file path (optional, uses env vars by default)
|
||||
|
||||
--no-onboard
|
||||
Skip first-run onboarding check
|
||||
|
||||
-h, --help
|
||||
Print help (see a summary with '-h')
|
||||
|
||||
-V, --version
|
||||
Print version
|
||||
@@ -23,6 +23,7 @@ Commands:
|
||||
service Manage OS service
|
||||
skills Manage skills
|
||||
doctor Run diagnostics
|
||||
logs View and manage gateway logs
|
||||
status Show system status
|
||||
completion Generate completions
|
||||
help Print this message or the help of the given subcommand(s)
|
||||
|
||||
+381
-34
@@ -91,11 +91,28 @@ pub struct SignalConfig {
|
||||
}
|
||||
|
||||
impl ChannelsConfig {
|
||||
pub(crate) fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
let http = if optional_env("HTTP_PORT")?.is_some() || optional_env("HTTP_HOST")?.is_some() {
|
||||
/// Resolve channels config following `env > settings > default` for every field.
|
||||
pub(crate) fn resolve(settings: &Settings, tunnel_enabled: bool) -> Result<Self, ConfigError> {
|
||||
let cs = &settings.channels;
|
||||
|
||||
// --- HTTP webhook ---
|
||||
// HTTP is enabled when env vars are set OR settings has it enabled.
|
||||
let http_enabled_by_env =
|
||||
optional_env("HTTP_PORT")?.is_some() || optional_env("HTTP_HOST")?.is_some();
|
||||
// When a tunnel is configured, default to loopback since external
|
||||
// traffic arrives through the tunnel. Without a tunnel the webhook
|
||||
// server needs to accept connections from the network directly.
|
||||
let default_host = if tunnel_enabled {
|
||||
"127.0.0.1"
|
||||
} else {
|
||||
"0.0.0.0"
|
||||
};
|
||||
let http = if http_enabled_by_env || cs.http_enabled {
|
||||
Some(HttpConfig {
|
||||
host: optional_env("HTTP_HOST")?.unwrap_or_else(|| "0.0.0.0".to_string()),
|
||||
port: parse_optional_env("HTTP_PORT", 8080)?,
|
||||
host: optional_env("HTTP_HOST")?
|
||||
.or_else(|| cs.http_host.clone())
|
||||
.unwrap_or_else(|| default_host.to_string()),
|
||||
port: parse_optional_env("HTTP_PORT", cs.http_port.unwrap_or(8080))?,
|
||||
webhook_secret: optional_env("HTTP_WEBHOOK_SECRET")?.map(SecretString::from),
|
||||
user_id: optional_env("HTTP_USER_ID")?.unwrap_or_else(|| "http".to_string()),
|
||||
})
|
||||
@@ -103,42 +120,58 @@ impl ChannelsConfig {
|
||||
None
|
||||
};
|
||||
|
||||
let gateway_enabled = parse_bool_env("GATEWAY_ENABLED", true)?;
|
||||
// --- Web gateway ---
|
||||
let gateway_enabled = parse_bool_env("GATEWAY_ENABLED", cs.gateway_enabled)?;
|
||||
let gateway = if gateway_enabled {
|
||||
Some(GatewayConfig {
|
||||
host: optional_env("GATEWAY_HOST")?.unwrap_or_else(|| "127.0.0.1".to_string()),
|
||||
port: parse_optional_env("GATEWAY_PORT", 3000)?,
|
||||
auth_token: optional_env("GATEWAY_AUTH_TOKEN")?,
|
||||
user_id: optional_env("GATEWAY_USER_ID")?.unwrap_or_else(|| "default".to_string()),
|
||||
host: optional_env("GATEWAY_HOST")?
|
||||
.or_else(|| cs.gateway_host.clone())
|
||||
.unwrap_or_else(|| "127.0.0.1".to_string()),
|
||||
port: parse_optional_env(
|
||||
"GATEWAY_PORT",
|
||||
cs.gateway_port.unwrap_or(DEFAULT_GATEWAY_PORT),
|
||||
)?,
|
||||
auth_token: optional_env("GATEWAY_AUTH_TOKEN")?
|
||||
.or_else(|| cs.gateway_auth_token.clone()),
|
||||
user_id: optional_env("GATEWAY_USER_ID")?
|
||||
.or_else(|| cs.gateway_user_id.clone())
|
||||
.unwrap_or_else(|| "default".to_string()),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let signal = if let Some(http_url) = optional_env("SIGNAL_HTTP_URL")? {
|
||||
let account = optional_env("SIGNAL_ACCOUNT")?.ok_or(ConfigError::InvalidValue {
|
||||
key: "SIGNAL_ACCOUNT".to_string(),
|
||||
message: "SIGNAL_ACCOUNT is required when SIGNAL_HTTP_URL is set".to_string(),
|
||||
})?;
|
||||
let allow_from = match std::env::var_os("SIGNAL_ALLOW_FROM") {
|
||||
// --- Signal ---
|
||||
let signal_url = optional_env("SIGNAL_HTTP_URL")?.or_else(|| cs.signal_http_url.clone());
|
||||
let signal = if let Some(http_url) = signal_url {
|
||||
let account = optional_env("SIGNAL_ACCOUNT")?
|
||||
.or_else(|| cs.signal_account.clone())
|
||||
.ok_or(ConfigError::InvalidValue {
|
||||
key: "SIGNAL_ACCOUNT".to_string(),
|
||||
message: "SIGNAL_ACCOUNT is required when Signal is enabled".to_string(),
|
||||
})?;
|
||||
let allow_from_str =
|
||||
optional_env("SIGNAL_ALLOW_FROM")?.or_else(|| cs.signal_allow_from.clone());
|
||||
let allow_from = match allow_from_str {
|
||||
None => vec![account.clone()],
|
||||
Some(val) => {
|
||||
let s = val.to_string_lossy();
|
||||
s.split(',')
|
||||
.map(|e| e.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect()
|
||||
}
|
||||
Some(s) => s
|
||||
.split(',')
|
||||
.map(|e| e.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect(),
|
||||
};
|
||||
let dm_policy =
|
||||
optional_env("SIGNAL_DM_POLICY")?.unwrap_or_else(|| "pairing".to_string());
|
||||
let group_policy =
|
||||
optional_env("SIGNAL_GROUP_POLICY")?.unwrap_or_else(|| "allowlist".to_string());
|
||||
let dm_policy = optional_env("SIGNAL_DM_POLICY")?
|
||||
.or_else(|| cs.signal_dm_policy.clone())
|
||||
.unwrap_or_else(|| "pairing".to_string());
|
||||
let group_policy = optional_env("SIGNAL_GROUP_POLICY")?
|
||||
.or_else(|| cs.signal_group_policy.clone())
|
||||
.unwrap_or_else(|| "allowlist".to_string());
|
||||
Some(SignalConfig {
|
||||
http_url,
|
||||
account,
|
||||
allow_from,
|
||||
allow_from_groups: optional_env("SIGNAL_ALLOW_FROM_GROUPS")?
|
||||
.or_else(|| cs.signal_allow_from_groups.clone())
|
||||
.map(|s| {
|
||||
s.split(',')
|
||||
.map(|e| e.trim().to_string())
|
||||
@@ -149,6 +182,7 @@ impl ChannelsConfig {
|
||||
dm_policy,
|
||||
group_policy,
|
||||
group_allow_from: optional_env("SIGNAL_GROUP_ALLOW_FROM")?
|
||||
.or_else(|| cs.signal_group_allow_from.clone())
|
||||
.map(|s| {
|
||||
s.split(',')
|
||||
.map(|e| e.trim().to_string())
|
||||
@@ -167,9 +201,17 @@ impl ChannelsConfig {
|
||||
None
|
||||
};
|
||||
|
||||
let cli_enabled = optional_env("CLI_ENABLED")?
|
||||
.map(|s| s.to_lowercase() != "false" && s != "0")
|
||||
.unwrap_or(true);
|
||||
// --- CLI ---
|
||||
let cli_enabled = parse_bool_env("CLI_ENABLED", cs.cli_enabled)?;
|
||||
|
||||
// --- WASM channels ---
|
||||
let wasm_channels_dir = optional_env("WASM_CHANNELS_DIR")?
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| cs.wasm_channels_dir.clone())
|
||||
.unwrap_or_else(default_channels_dir);
|
||||
|
||||
let wasm_channels_enabled =
|
||||
parse_bool_env("WASM_CHANNELS_ENABLED", cs.wasm_channels_enabled)?;
|
||||
|
||||
Ok(Self {
|
||||
cli: CliConfig {
|
||||
@@ -178,12 +220,10 @@ impl ChannelsConfig {
|
||||
http,
|
||||
gateway,
|
||||
signal,
|
||||
wasm_channels_dir: optional_env("WASM_CHANNELS_DIR")?
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(default_channels_dir),
|
||||
wasm_channels_enabled: parse_bool_env("WASM_CHANNELS_ENABLED", true)?,
|
||||
wasm_channels_dir,
|
||||
wasm_channels_enabled,
|
||||
wasm_channel_owner_ids: {
|
||||
let mut ids = settings.channels.wasm_channel_owner_ids.clone();
|
||||
let mut ids = cs.wasm_channel_owner_ids.clone();
|
||||
// Backwards compat: TELEGRAM_OWNER_ID env var
|
||||
if let Some(id_str) = optional_env("TELEGRAM_OWNER_ID")? {
|
||||
let id: i64 = id_str.parse().map_err(|e: std::num::ParseIntError| {
|
||||
@@ -200,6 +240,10 @@ impl ChannelsConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Default gateway port — used both in `resolve()` and as the fallback in
|
||||
/// other modules that need to construct a gateway URL.
|
||||
pub const DEFAULT_GATEWAY_PORT: u16 = 3000;
|
||||
|
||||
/// Get the default channels directory (~/.ironclaw/channels/).
|
||||
fn default_channels_dir() -> PathBuf {
|
||||
ironclaw_base_dir().join("channels")
|
||||
@@ -354,6 +398,69 @@ mod tests {
|
||||
assert!(!cfg.wasm_channels_enabled);
|
||||
}
|
||||
|
||||
/// When a tunnel is active and HTTP_HOST is not explicitly set, the
|
||||
/// webhook server should default to loopback to avoid unnecessary exposure.
|
||||
#[test]
|
||||
fn http_host_defaults_to_loopback_with_tunnel() {
|
||||
// Set HTTP_PORT to trigger HttpConfig creation, but leave HTTP_HOST unset
|
||||
// so the default kicks in.
|
||||
unsafe {
|
||||
std::env::set_var("HTTP_PORT", "9999");
|
||||
std::env::remove_var("HTTP_HOST");
|
||||
}
|
||||
let settings = crate::settings::Settings::default();
|
||||
let cfg = ChannelsConfig::resolve(&settings, true).unwrap();
|
||||
unsafe {
|
||||
std::env::remove_var("HTTP_PORT");
|
||||
}
|
||||
let http = cfg.http.expect("HttpConfig should be present");
|
||||
assert_eq!(
|
||||
http.host, "127.0.0.1",
|
||||
"tunnel active should default to loopback"
|
||||
);
|
||||
assert_eq!(http.port, 9999);
|
||||
}
|
||||
|
||||
/// Without a tunnel, the webhook server defaults to 0.0.0.0 so external
|
||||
/// services can reach it directly.
|
||||
#[test]
|
||||
fn http_host_defaults_to_all_interfaces_without_tunnel() {
|
||||
unsafe {
|
||||
std::env::set_var("HTTP_PORT", "9998");
|
||||
std::env::remove_var("HTTP_HOST");
|
||||
}
|
||||
let settings = crate::settings::Settings::default();
|
||||
let cfg = ChannelsConfig::resolve(&settings, false).unwrap();
|
||||
unsafe {
|
||||
std::env::remove_var("HTTP_PORT");
|
||||
}
|
||||
let http = cfg.http.expect("HttpConfig should be present");
|
||||
assert_eq!(
|
||||
http.host, "0.0.0.0",
|
||||
"no tunnel should default to all interfaces"
|
||||
);
|
||||
}
|
||||
|
||||
/// An explicit HTTP_HOST always wins regardless of tunnel state.
|
||||
#[test]
|
||||
fn explicit_http_host_overrides_tunnel_default() {
|
||||
unsafe {
|
||||
std::env::set_var("HTTP_PORT", "9997");
|
||||
std::env::set_var("HTTP_HOST", "192.168.1.50");
|
||||
}
|
||||
let settings = crate::settings::Settings::default();
|
||||
let cfg = ChannelsConfig::resolve(&settings, true).unwrap();
|
||||
unsafe {
|
||||
std::env::remove_var("HTTP_PORT");
|
||||
std::env::remove_var("HTTP_HOST");
|
||||
}
|
||||
let http = cfg.http.expect("HttpConfig should be present");
|
||||
assert_eq!(
|
||||
http.host, "192.168.1.50",
|
||||
"explicit host should override tunnel default"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_channels_dir_ends_with_channels() {
|
||||
let dir = default_channels_dir();
|
||||
@@ -362,4 +469,244 @@ mod tests {
|
||||
"expected path ending in 'channels', got: {dir:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_gateway_port_constant() {
|
||||
assert_eq!(DEFAULT_GATEWAY_PORT, 3000);
|
||||
}
|
||||
|
||||
/// With default settings and no env vars, gateway should use defaults.
|
||||
#[test]
|
||||
fn resolve_gateway_defaults_from_settings() {
|
||||
let _lock = crate::config::helpers::ENV_MUTEX.lock();
|
||||
// Clear env vars that would interfere
|
||||
unsafe {
|
||||
std::env::remove_var("GATEWAY_ENABLED");
|
||||
std::env::remove_var("GATEWAY_HOST");
|
||||
std::env::remove_var("GATEWAY_PORT");
|
||||
std::env::remove_var("GATEWAY_AUTH_TOKEN");
|
||||
std::env::remove_var("GATEWAY_USER_ID");
|
||||
std::env::remove_var("HTTP_PORT");
|
||||
std::env::remove_var("HTTP_HOST");
|
||||
std::env::remove_var("SIGNAL_HTTP_URL");
|
||||
std::env::remove_var("CLI_ENABLED");
|
||||
std::env::remove_var("WASM_CHANNELS_DIR");
|
||||
std::env::remove_var("WASM_CHANNELS_ENABLED");
|
||||
std::env::remove_var("TELEGRAM_OWNER_ID");
|
||||
}
|
||||
|
||||
let settings = crate::settings::Settings::default();
|
||||
let cfg = ChannelsConfig::resolve(&settings, false).unwrap();
|
||||
|
||||
let gw = cfg.gateway.expect("gateway should be enabled by default");
|
||||
assert_eq!(gw.host, "127.0.0.1");
|
||||
assert_eq!(gw.port, DEFAULT_GATEWAY_PORT);
|
||||
assert!(gw.auth_token.is_none());
|
||||
assert_eq!(gw.user_id, "default");
|
||||
}
|
||||
|
||||
/// Settings values should be used when no env vars are set.
|
||||
#[test]
|
||||
fn resolve_gateway_from_settings() {
|
||||
let _lock = crate::config::helpers::ENV_MUTEX.lock();
|
||||
unsafe {
|
||||
std::env::remove_var("GATEWAY_ENABLED");
|
||||
std::env::remove_var("GATEWAY_HOST");
|
||||
std::env::remove_var("GATEWAY_PORT");
|
||||
std::env::remove_var("GATEWAY_AUTH_TOKEN");
|
||||
std::env::remove_var("GATEWAY_USER_ID");
|
||||
std::env::remove_var("HTTP_PORT");
|
||||
std::env::remove_var("HTTP_HOST");
|
||||
std::env::remove_var("SIGNAL_HTTP_URL");
|
||||
std::env::remove_var("CLI_ENABLED");
|
||||
std::env::remove_var("WASM_CHANNELS_DIR");
|
||||
std::env::remove_var("WASM_CHANNELS_ENABLED");
|
||||
std::env::remove_var("TELEGRAM_OWNER_ID");
|
||||
}
|
||||
|
||||
let mut settings = crate::settings::Settings::default();
|
||||
settings.channels.gateway_port = Some(4000);
|
||||
settings.channels.gateway_host = Some("0.0.0.0".to_string());
|
||||
settings.channels.gateway_auth_token = Some("db-token-123".to_string());
|
||||
settings.channels.gateway_user_id = Some("myuser".to_string());
|
||||
|
||||
let cfg = ChannelsConfig::resolve(&settings, false).unwrap();
|
||||
let gw = cfg.gateway.expect("gateway should be enabled");
|
||||
assert_eq!(gw.port, 4000);
|
||||
assert_eq!(gw.host, "0.0.0.0");
|
||||
assert_eq!(gw.auth_token.as_deref(), Some("db-token-123"));
|
||||
assert_eq!(gw.user_id, "myuser");
|
||||
}
|
||||
|
||||
/// Env vars should override settings values.
|
||||
#[test]
|
||||
fn resolve_env_overrides_settings() {
|
||||
let _lock = crate::config::helpers::ENV_MUTEX.lock();
|
||||
unsafe {
|
||||
std::env::set_var("GATEWAY_PORT", "5000");
|
||||
std::env::set_var("GATEWAY_HOST", "10.0.0.1");
|
||||
std::env::set_var("GATEWAY_AUTH_TOKEN", "env-token");
|
||||
std::env::remove_var("GATEWAY_ENABLED");
|
||||
std::env::remove_var("GATEWAY_USER_ID");
|
||||
std::env::remove_var("HTTP_PORT");
|
||||
std::env::remove_var("HTTP_HOST");
|
||||
std::env::remove_var("SIGNAL_HTTP_URL");
|
||||
std::env::remove_var("CLI_ENABLED");
|
||||
std::env::remove_var("WASM_CHANNELS_DIR");
|
||||
std::env::remove_var("WASM_CHANNELS_ENABLED");
|
||||
std::env::remove_var("TELEGRAM_OWNER_ID");
|
||||
}
|
||||
|
||||
let mut settings = crate::settings::Settings::default();
|
||||
settings.channels.gateway_port = Some(4000);
|
||||
settings.channels.gateway_host = Some("0.0.0.0".to_string());
|
||||
settings.channels.gateway_auth_token = Some("db-token".to_string());
|
||||
|
||||
let cfg = ChannelsConfig::resolve(&settings, false).unwrap();
|
||||
let gw = cfg.gateway.expect("gateway should be enabled");
|
||||
assert_eq!(gw.port, 5000, "env should override settings");
|
||||
assert_eq!(gw.host, "10.0.0.1", "env should override settings");
|
||||
assert_eq!(
|
||||
gw.auth_token.as_deref(),
|
||||
Some("env-token"),
|
||||
"env should override settings"
|
||||
);
|
||||
|
||||
// Cleanup
|
||||
unsafe {
|
||||
std::env::remove_var("GATEWAY_PORT");
|
||||
std::env::remove_var("GATEWAY_HOST");
|
||||
std::env::remove_var("GATEWAY_AUTH_TOKEN");
|
||||
}
|
||||
}
|
||||
|
||||
/// CLI enabled should fall back to settings.
|
||||
#[test]
|
||||
fn resolve_cli_enabled_from_settings() {
|
||||
let _lock = crate::config::helpers::ENV_MUTEX.lock();
|
||||
unsafe {
|
||||
std::env::remove_var("CLI_ENABLED");
|
||||
std::env::remove_var("GATEWAY_ENABLED");
|
||||
std::env::remove_var("GATEWAY_HOST");
|
||||
std::env::remove_var("GATEWAY_PORT");
|
||||
std::env::remove_var("GATEWAY_AUTH_TOKEN");
|
||||
std::env::remove_var("GATEWAY_USER_ID");
|
||||
std::env::remove_var("HTTP_PORT");
|
||||
std::env::remove_var("HTTP_HOST");
|
||||
std::env::remove_var("SIGNAL_HTTP_URL");
|
||||
std::env::remove_var("WASM_CHANNELS_DIR");
|
||||
std::env::remove_var("WASM_CHANNELS_ENABLED");
|
||||
std::env::remove_var("TELEGRAM_OWNER_ID");
|
||||
}
|
||||
|
||||
let mut settings = crate::settings::Settings::default();
|
||||
settings.channels.cli_enabled = false;
|
||||
|
||||
let cfg = ChannelsConfig::resolve(&settings, false).unwrap();
|
||||
assert!(!cfg.cli.enabled, "settings should disable CLI");
|
||||
}
|
||||
|
||||
/// HTTP channel should activate when settings has it enabled.
|
||||
#[test]
|
||||
fn resolve_http_from_settings() {
|
||||
let _lock = crate::config::helpers::ENV_MUTEX.lock();
|
||||
unsafe {
|
||||
std::env::remove_var("HTTP_PORT");
|
||||
std::env::remove_var("HTTP_HOST");
|
||||
std::env::remove_var("HTTP_WEBHOOK_SECRET");
|
||||
std::env::remove_var("HTTP_USER_ID");
|
||||
std::env::remove_var("GATEWAY_ENABLED");
|
||||
std::env::remove_var("GATEWAY_HOST");
|
||||
std::env::remove_var("GATEWAY_PORT");
|
||||
std::env::remove_var("GATEWAY_AUTH_TOKEN");
|
||||
std::env::remove_var("GATEWAY_USER_ID");
|
||||
std::env::remove_var("SIGNAL_HTTP_URL");
|
||||
std::env::remove_var("CLI_ENABLED");
|
||||
std::env::remove_var("WASM_CHANNELS_DIR");
|
||||
std::env::remove_var("WASM_CHANNELS_ENABLED");
|
||||
std::env::remove_var("TELEGRAM_OWNER_ID");
|
||||
}
|
||||
|
||||
let mut settings = crate::settings::Settings::default();
|
||||
settings.channels.http_enabled = true;
|
||||
settings.channels.http_port = Some(9090);
|
||||
settings.channels.http_host = Some("10.0.0.1".to_string());
|
||||
|
||||
let cfg = ChannelsConfig::resolve(&settings, false).unwrap();
|
||||
let http = cfg.http.expect("HTTP should be enabled from settings");
|
||||
assert_eq!(http.port, 9090);
|
||||
assert_eq!(http.host, "10.0.0.1");
|
||||
}
|
||||
|
||||
/// Settings round-trip through DB map for new gateway fields.
|
||||
#[test]
|
||||
fn settings_gateway_fields_db_roundtrip() {
|
||||
let mut settings = crate::settings::Settings::default();
|
||||
settings.channels.gateway_port = Some(4000);
|
||||
settings.channels.gateway_host = Some("0.0.0.0".to_string());
|
||||
settings.channels.gateway_auth_token = Some("tok-abc".to_string());
|
||||
settings.channels.gateway_user_id = Some("myuser".to_string());
|
||||
settings.channels.cli_enabled = false;
|
||||
|
||||
let map = settings.to_db_map();
|
||||
let restored = crate::settings::Settings::from_db_map(&map);
|
||||
|
||||
assert_eq!(restored.channels.gateway_port, Some(4000));
|
||||
assert_eq!(restored.channels.gateway_host.as_deref(), Some("0.0.0.0"));
|
||||
assert_eq!(
|
||||
restored.channels.gateway_auth_token.as_deref(),
|
||||
Some("tok-abc")
|
||||
);
|
||||
assert_eq!(restored.channels.gateway_user_id.as_deref(), Some("myuser"));
|
||||
assert!(!restored.channels.cli_enabled);
|
||||
}
|
||||
|
||||
/// Invalid boolean env values must produce errors, not silently degrade.
|
||||
#[test]
|
||||
fn resolve_rejects_invalid_bool_env() {
|
||||
let _lock = crate::config::helpers::ENV_MUTEX.lock();
|
||||
let settings = crate::settings::Settings::default();
|
||||
|
||||
// GATEWAY_ENABLED=maybe should error
|
||||
unsafe {
|
||||
std::env::set_var("GATEWAY_ENABLED", "maybe");
|
||||
std::env::remove_var("HTTP_PORT");
|
||||
std::env::remove_var("HTTP_HOST");
|
||||
std::env::remove_var("SIGNAL_HTTP_URL");
|
||||
std::env::remove_var("CLI_ENABLED");
|
||||
std::env::remove_var("WASM_CHANNELS_ENABLED");
|
||||
std::env::remove_var("GATEWAY_PORT");
|
||||
std::env::remove_var("GATEWAY_HOST");
|
||||
std::env::remove_var("GATEWAY_AUTH_TOKEN");
|
||||
std::env::remove_var("GATEWAY_USER_ID");
|
||||
std::env::remove_var("WASM_CHANNELS_DIR");
|
||||
std::env::remove_var("TELEGRAM_OWNER_ID");
|
||||
}
|
||||
let result = ChannelsConfig::resolve(&settings, false);
|
||||
assert!(result.is_err(), "GATEWAY_ENABLED=maybe should be rejected");
|
||||
|
||||
// CLI_ENABLED=on should error
|
||||
unsafe {
|
||||
std::env::remove_var("GATEWAY_ENABLED");
|
||||
std::env::set_var("CLI_ENABLED", "on");
|
||||
}
|
||||
let result = ChannelsConfig::resolve(&settings, false);
|
||||
assert!(result.is_err(), "CLI_ENABLED=on should be rejected");
|
||||
|
||||
// WASM_CHANNELS_ENABLED=yes should error
|
||||
unsafe {
|
||||
std::env::remove_var("CLI_ENABLED");
|
||||
std::env::set_var("WASM_CHANNELS_ENABLED", "yes");
|
||||
}
|
||||
let result = ChannelsConfig::resolve(&settings, false);
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"WASM_CHANNELS_ENABLED=yes should be rejected"
|
||||
);
|
||||
|
||||
// Cleanup
|
||||
unsafe {
|
||||
std::env::remove_var("WASM_CHANNELS_ENABLED");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,6 +170,40 @@ impl DatabaseConfig {
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a config from a raw PostgreSQL URL (for wizard/testing).
|
||||
pub fn from_postgres_url(url: &str, pool_size: usize) -> Self {
|
||||
Self {
|
||||
backend: DatabaseBackend::Postgres,
|
||||
url: SecretString::from(url.to_string()),
|
||||
pool_size,
|
||||
ssl_mode: SslMode::from_env(),
|
||||
libsql_path: None,
|
||||
libsql_url: None,
|
||||
libsql_auth_token: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a config for a libSQL database (for wizard/testing).
|
||||
///
|
||||
/// Empty strings for `turso_url` and `turso_token` are treated as `None`.
|
||||
pub fn from_libsql_path(
|
||||
path: &str,
|
||||
turso_url: Option<&str>,
|
||||
turso_token: Option<&str>,
|
||||
) -> Self {
|
||||
let turso_url = turso_url.filter(|s| !s.is_empty());
|
||||
let turso_token = turso_token.filter(|s| !s.is_empty());
|
||||
Self {
|
||||
backend: DatabaseBackend::LibSql,
|
||||
url: SecretString::from("unused://libsql".to_string()),
|
||||
pool_size: 1,
|
||||
ssl_mode: SslMode::default(),
|
||||
libsql_path: Some(PathBuf::from(path)),
|
||||
libsql_url: turso_url.map(String::from),
|
||||
libsql_auth_token: turso_token.map(|t| SecretString::from(t.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the database URL (exposes the secret).
|
||||
pub fn url(&self) -> &str {
|
||||
self.url.expose_secret()
|
||||
|
||||
+9
-3
@@ -34,7 +34,9 @@ use crate::settings::Settings;
|
||||
// Re-export all public types so `crate::config::FooConfig` continues to work.
|
||||
pub use self::agent::AgentConfig;
|
||||
pub use self::builder::BuilderModeConfig;
|
||||
pub use self::channels::{ChannelsConfig, CliConfig, GatewayConfig, HttpConfig, SignalConfig};
|
||||
pub use self::channels::{
|
||||
ChannelsConfig, CliConfig, DEFAULT_GATEWAY_PORT, GatewayConfig, HttpConfig, SignalConfig,
|
||||
};
|
||||
pub use self::database::{DatabaseBackend, DatabaseConfig, SslMode, default_libsql_path};
|
||||
pub use self::embeddings::EmbeddingsConfig;
|
||||
pub use self::heartbeat::HeartbeatConfig;
|
||||
@@ -304,12 +306,16 @@ impl Config {
|
||||
|
||||
/// Build config from settings (shared by from_env and from_db).
|
||||
async fn build(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
// Resolve tunnel first so channels can default to loopback when a
|
||||
// tunnel handles external exposure (no need to bind 0.0.0.0).
|
||||
let tunnel = TunnelConfig::resolve(settings)?;
|
||||
|
||||
Ok(Self {
|
||||
database: DatabaseConfig::resolve()?,
|
||||
llm: LlmConfig::resolve(settings)?,
|
||||
embeddings: EmbeddingsConfig::resolve(settings)?,
|
||||
tunnel: TunnelConfig::resolve(settings)?,
|
||||
channels: ChannelsConfig::resolve(settings)?,
|
||||
channels: ChannelsConfig::resolve(settings, tunnel.is_enabled())?,
|
||||
tunnel,
|
||||
agent: AgentConfig::resolve(settings)?,
|
||||
safety: resolve_safety_config()?,
|
||||
wasm: WasmConfig::resolve()?,
|
||||
|
||||
@@ -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,
|
||||
|
||||
+144
-11
@@ -104,7 +104,7 @@ pub async fn connect_with_handles(
|
||||
Ok((Arc::new(backend) as Arc<dyn Database>, handles))
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
_ => {
|
||||
crate::config::DatabaseBackend::Postgres => {
|
||||
let pg = postgres::PgBackend::new(config)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
|
||||
@@ -115,10 +115,11 @@ pub async fn connect_with_handles(
|
||||
|
||||
Ok((Arc::new(pg) as Arc<dyn Database>, handles))
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
_ => Err(DatabaseError::Pool(
|
||||
"No database backend available. Enable 'postgres' or 'libsql' feature.".to_string(),
|
||||
)),
|
||||
#[allow(unreachable_patterns)]
|
||||
_ => Err(DatabaseError::Pool(format!(
|
||||
"Database backend '{}' is not available. Rebuild with the appropriate feature flag.",
|
||||
config.backend
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -161,7 +162,7 @@ pub async fn create_secrets_store(
|
||||
)))
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
_ => {
|
||||
crate::config::DatabaseBackend::Postgres => {
|
||||
let pg = postgres::PgBackend::new(config)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
|
||||
@@ -172,14 +173,142 @@ pub async fn create_secrets_store(
|
||||
crypto,
|
||||
)))
|
||||
}
|
||||
#[cfg(not(feature = "postgres"))]
|
||||
_ => Err(DatabaseError::Pool(
|
||||
"No database backend available for secrets. Enable 'postgres' or 'libsql' feature."
|
||||
.to_string(),
|
||||
)),
|
||||
#[allow(unreachable_patterns)]
|
||||
_ => Err(DatabaseError::Pool(format!(
|
||||
"Database backend '{}' is not available for secrets. Rebuild with the appropriate feature flag.",
|
||||
config.backend
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== Wizard / testing helpers ====================
|
||||
|
||||
/// Connect to the database WITHOUT running migrations, validating
|
||||
/// prerequisites when applicable (PostgreSQL version, pgvector).
|
||||
///
|
||||
/// Returns both the `Database` trait object and backend-specific handles.
|
||||
/// Used by the wizard to test connectivity before committing — call
|
||||
/// [`Database::run_migrations`] on the returned trait object when ready.
|
||||
pub async fn connect_without_migrations(
|
||||
config: &crate::config::DatabaseConfig,
|
||||
) -> Result<(Arc<dyn Database>, DatabaseHandles), DatabaseError> {
|
||||
let mut handles = DatabaseHandles::default();
|
||||
|
||||
match config.backend {
|
||||
#[cfg(feature = "libsql")]
|
||||
crate::config::DatabaseBackend::LibSql => {
|
||||
use secrecy::ExposeSecret as _;
|
||||
|
||||
let default_path = crate::config::default_libsql_path();
|
||||
let db_path = config.libsql_path.as_deref().unwrap_or(&default_path);
|
||||
|
||||
let backend = if let Some(ref url) = config.libsql_url {
|
||||
let token = config.libsql_auth_token.as_ref().ok_or_else(|| {
|
||||
DatabaseError::Pool(
|
||||
"LIBSQL_AUTH_TOKEN required when LIBSQL_URL is set".to_string(),
|
||||
)
|
||||
})?;
|
||||
libsql::LibSqlBackend::new_remote_replica(db_path, url, token.expose_secret())
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?
|
||||
} else {
|
||||
libsql::LibSqlBackend::new_local(db_path)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?
|
||||
};
|
||||
|
||||
handles.libsql_db = Some(backend.shared_db());
|
||||
|
||||
Ok((Arc::new(backend) as Arc<dyn Database>, handles))
|
||||
}
|
||||
#[cfg(feature = "postgres")]
|
||||
crate::config::DatabaseBackend::Postgres => {
|
||||
let pg = postgres::PgBackend::new(config)
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
|
||||
|
||||
handles.pg_pool = Some(pg.pool());
|
||||
|
||||
// Validate PostgreSQL prerequisites (version, pgvector)
|
||||
validate_postgres(&pg.pool()).await?;
|
||||
|
||||
Ok((Arc::new(pg) as Arc<dyn Database>, handles))
|
||||
}
|
||||
#[allow(unreachable_patterns)]
|
||||
_ => Err(DatabaseError::Pool(format!(
|
||||
"Database backend '{}' is not available. Rebuild with the appropriate feature flag.",
|
||||
config.backend
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate PostgreSQL prerequisites (version >= 15, pgvector available).
|
||||
///
|
||||
/// Returns `Ok(())` if all prerequisites are met, or a `DatabaseError`
|
||||
/// with a user-facing message describing the issue.
|
||||
#[cfg(feature = "postgres")]
|
||||
async fn validate_postgres(pool: &deadpool_postgres::Pool) -> Result<(), DatabaseError> {
|
||||
let client = pool
|
||||
.get()
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Pool(format!("Failed to connect: {}", e)))?;
|
||||
|
||||
// Check PostgreSQL server version (need 15+ for pgvector).
|
||||
let version_row = client
|
||||
.query_one("SHOW server_version", &[])
|
||||
.await
|
||||
.map_err(|e| DatabaseError::Query(format!("Failed to query server version: {}", e)))?;
|
||||
let version_str: &str = version_row.get(0);
|
||||
let major_version = version_str
|
||||
.split('.')
|
||||
.next()
|
||||
.and_then(|v| v.parse::<u32>().ok())
|
||||
.ok_or_else(|| {
|
||||
DatabaseError::Pool(format!(
|
||||
"Could not parse PostgreSQL version from '{}'. \
|
||||
Expected a numeric major version (e.g., '15.2').",
|
||||
version_str
|
||||
))
|
||||
})?;
|
||||
|
||||
const MIN_PG_MAJOR_VERSION: u32 = 15;
|
||||
|
||||
if major_version < MIN_PG_MAJOR_VERSION {
|
||||
return Err(DatabaseError::Pool(format!(
|
||||
"PostgreSQL {} detected. IronClaw requires PostgreSQL {} or later \
|
||||
for pgvector support.\n\
|
||||
Upgrade: https://www.postgresql.org/download/",
|
||||
version_str, MIN_PG_MAJOR_VERSION
|
||||
)));
|
||||
}
|
||||
|
||||
// Check if pgvector extension is available.
|
||||
let pgvector_row = client
|
||||
.query_opt(
|
||||
"SELECT 1 FROM pg_available_extensions WHERE name = 'vector'",
|
||||
&[],
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
DatabaseError::Query(format!("Failed to check pgvector availability: {}", e))
|
||||
})?;
|
||||
|
||||
if pgvector_row.is_none() {
|
||||
return Err(DatabaseError::Pool(format!(
|
||||
"pgvector extension not found on your PostgreSQL server.\n\n\
|
||||
Install it:\n \
|
||||
macOS: brew install pgvector\n \
|
||||
Ubuntu: apt install postgresql-{0}-pgvector\n \
|
||||
Docker: use the pgvector/pgvector:pg{0} image\n \
|
||||
Source: https://github.com/pgvector/pgvector#installation\n\n\
|
||||
Then restart PostgreSQL and re-run: ironclaw onboard",
|
||||
major_version
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ==================== Sub-traits ====================
|
||||
//
|
||||
// Each sub-trait groups related persistence methods. The `Database` supertrait
|
||||
@@ -387,6 +516,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,
|
||||
|
||||
@@ -205,7 +205,8 @@ fn extract_rtf(data: &[u8]) -> Result<String, String> {
|
||||
let mut word = String::new();
|
||||
while let Some(&next) = chars.peek() {
|
||||
if next.is_ascii_alphabetic() {
|
||||
word.push(chars.next().unwrap());
|
||||
chars.next();
|
||||
word.push(next);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
|
||||
+177
-21
@@ -248,12 +248,14 @@ impl ExtensionManager {
|
||||
self.tunnel_url
|
||||
.as_ref()
|
||||
.filter(|u| !u.is_empty())
|
||||
.and_then(|raw| url::Url::parse(raw).ok())
|
||||
.and_then(|u| u.host_str().map(String::from))
|
||||
.filter(|host| !oauth_defaults::is_loopback_host(host))
|
||||
.map(|_| {
|
||||
let base = self.tunnel_url.as_ref().unwrap().trim_end_matches('/');
|
||||
format!("{}/oauth/callback", base)
|
||||
.and_then(|raw| {
|
||||
let url = url::Url::parse(raw).ok()?;
|
||||
let host = url.host_str().map(String::from)?;
|
||||
if oauth_defaults::is_loopback_host(&host) {
|
||||
return None;
|
||||
}
|
||||
let base = raw.trim_end_matches('/');
|
||||
Some(format!("{}/oauth/callback", base))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -304,6 +306,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
|
||||
@@ -1281,8 +1311,12 @@ impl ExtensionManager {
|
||||
match fallback_decision(&primary_result, &entry.fallback_source) {
|
||||
FallbackDecision::Return => primary_result,
|
||||
FallbackDecision::TryFallback => {
|
||||
let primary_err = primary_result.unwrap_err();
|
||||
let fallback = entry.fallback_source.as_ref().unwrap();
|
||||
// TryFallback guarantees primary is Err and fallback_source is Some.
|
||||
let (primary_err, fallback) = match (primary_result, entry.fallback_source.as_ref())
|
||||
{
|
||||
(Err(e), Some(f)) => (e, f),
|
||||
(other, _) => return other,
|
||||
};
|
||||
tracing::info!(
|
||||
extension = %entry.name,
|
||||
primary_error = %primary_err,
|
||||
@@ -2830,9 +2864,16 @@ impl ExtensionManager {
|
||||
// Try to list and create tools.
|
||||
// A 401/auth error means the server requires OAuth — surface as
|
||||
// AuthRequired so the activate handler triggers the OAuth flow.
|
||||
// Some servers (e.g. GitHub MCP) return 400 with "Authorization header
|
||||
// is badly formatted" instead of 401 when auth is missing or invalid.
|
||||
let mcp_tools = client.list_tools().await.map_err(|e| {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("requires authentication") || msg.contains("401") {
|
||||
let msg_lower = msg.to_ascii_lowercase();
|
||||
if msg_lower.contains("requires authentication")
|
||||
|| msg.contains("401")
|
||||
|| (msg.contains("400")
|
||||
&& (msg_lower.contains("authorization") || msg_lower.contains("authenticate")))
|
||||
{
|
||||
ExtensionError::AuthRequired
|
||||
} else {
|
||||
ExtensionError::ActivationFailed(msg)
|
||||
@@ -2980,13 +3021,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 +3031,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 +3101,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));
|
||||
}
|
||||
|
||||
@@ -3417,7 +3451,8 @@ impl ExtensionManager {
|
||||
.or_else(|| relay_config.callback_url.clone())
|
||||
.unwrap_or_else(|| {
|
||||
let host = std::env::var("GATEWAY_HOST").unwrap_or_else(|_| "127.0.0.1".into());
|
||||
let port = std::env::var("GATEWAY_PORT").unwrap_or_else(|_| "3001".into());
|
||||
let port = std::env::var("GATEWAY_PORT")
|
||||
.unwrap_or_else(|_| crate::config::DEFAULT_GATEWAY_PORT.to_string());
|
||||
format!("http://{}:{}", host, port)
|
||||
});
|
||||
|
||||
@@ -3816,11 +3851,12 @@ impl ExtensionManager {
|
||||
secret_name, name
|
||||
)));
|
||||
}
|
||||
if secret_value.trim().is_empty() {
|
||||
let trimmed_value = secret_value.trim();
|
||||
if trimmed_value.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let params =
|
||||
CreateSecretParams::new(secret_name, secret_value).with_provider(name.to_string());
|
||||
CreateSecretParams::new(secret_name, trimmed_value).with_provider(name.to_string());
|
||||
self.secrets
|
||||
.create(&self.user_id, params)
|
||||
.await
|
||||
@@ -4744,6 +4780,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]
|
||||
|
||||
+79
-226
@@ -232,198 +232,11 @@ pub fn builtin_entries() -> Vec<RegistryEntry> {
|
||||
}
|
||||
|
||||
/// Well-known extensions, with an optional relay URL for the channel-relay entry.
|
||||
///
|
||||
/// MCP server entries are loaded from `registry/mcp-servers/*.json` via the catalog
|
||||
/// system. Only runtime-dependent entries (like channel-relay) remain here.
|
||||
pub fn builtin_entries_with_relay(relay_url: Option<String>) -> Vec<RegistryEntry> {
|
||||
let mut entries = vec![
|
||||
// -- MCP Servers --
|
||||
RegistryEntry {
|
||||
name: "notion".to_string(),
|
||||
display_name: "Notion".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description: "Connect to Notion for reading and writing pages, databases, and comments"
|
||||
.to_string(),
|
||||
keywords: vec![
|
||||
"notes".into(),
|
||||
"wiki".into(),
|
||||
"docs".into(),
|
||||
"pages".into(),
|
||||
"database".into(),
|
||||
],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://mcp.notion.com/mcp".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "linear".to_string(),
|
||||
display_name: "Linear".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description:
|
||||
"Connect to Linear for issue tracking, project management, and team workflows"
|
||||
.to_string(),
|
||||
keywords: vec![
|
||||
"issues".into(),
|
||||
"tickets".into(),
|
||||
"project".into(),
|
||||
"tracking".into(),
|
||||
"bugs".into(),
|
||||
],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://mcp.linear.app/sse".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "github".to_string(),
|
||||
display_name: "GitHub".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description:
|
||||
"Connect to GitHub for repository management, issues, PRs, and code search"
|
||||
.to_string(),
|
||||
keywords: vec![
|
||||
"git".into(),
|
||||
"repos".into(),
|
||||
"code".into(),
|
||||
"pull-request".into(),
|
||||
"issues".into(),
|
||||
],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://api.githubcopilot.com/mcp/".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "slack-mcp".to_string(),
|
||||
display_name: "Slack MCP".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description:
|
||||
"Connect to Slack via MCP for messaging, channel management, and team communication"
|
||||
.to_string(),
|
||||
keywords: vec![
|
||||
"messaging".into(),
|
||||
"chat".into(),
|
||||
"channels".into(),
|
||||
"team".into(),
|
||||
"communication".into(),
|
||||
],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://mcp.slack.com".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "sentry".to_string(),
|
||||
display_name: "Sentry".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description:
|
||||
"Connect to Sentry for error tracking, performance monitoring, and debugging"
|
||||
.to_string(),
|
||||
keywords: vec![
|
||||
"errors".into(),
|
||||
"monitoring".into(),
|
||||
"debugging".into(),
|
||||
"crashes".into(),
|
||||
"performance".into(),
|
||||
],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://mcp.sentry.dev/mcp".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "stripe".to_string(),
|
||||
display_name: "Stripe".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description:
|
||||
"Connect to Stripe for payment processing, subscriptions, and financial data"
|
||||
.to_string(),
|
||||
keywords: vec![
|
||||
"payments".into(),
|
||||
"billing".into(),
|
||||
"subscriptions".into(),
|
||||
"invoices".into(),
|
||||
"finance".into(),
|
||||
],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://mcp.stripe.com".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "cloudflare".to_string(),
|
||||
display_name: "Cloudflare".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description:
|
||||
"Connect to Cloudflare for DNS, Workers, KV, and infrastructure management"
|
||||
.to_string(),
|
||||
keywords: vec![
|
||||
"cdn".into(),
|
||||
"dns".into(),
|
||||
"workers".into(),
|
||||
"hosting".into(),
|
||||
"infrastructure".into(),
|
||||
],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://mcp.cloudflare.com/mcp".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "asana".to_string(),
|
||||
display_name: "Asana".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description: "Connect to Asana for task management, projects, and team coordination"
|
||||
.to_string(),
|
||||
keywords: vec![
|
||||
"tasks".into(),
|
||||
"projects".into(),
|
||||
"management".into(),
|
||||
"team".into(),
|
||||
],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://mcp.asana.com/v2/mcp".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "intercom".to_string(),
|
||||
display_name: "Intercom".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description: "Connect to Intercom for customer messaging, support, and engagement"
|
||||
.to_string(),
|
||||
keywords: vec![
|
||||
"support".into(),
|
||||
"customers".into(),
|
||||
"messaging".into(),
|
||||
"chat".into(),
|
||||
"helpdesk".into(),
|
||||
],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://mcp.intercom.com/mcp".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
// WASM channels (telegram, slack, discord, whatsapp) come from the embedded
|
||||
// registry catalog (registry/channels/*.json) with WasmDownload URLs pointing
|
||||
// to GitHub release artifacts. See new_with_catalog() for merging.
|
||||
];
|
||||
let mut entries = vec![];
|
||||
|
||||
// Conditionally add channel-relay entries when relay URL is configured
|
||||
if let Some(relay_url) = relay_url {
|
||||
@@ -545,9 +358,21 @@ mod tests {
|
||||
assert_eq!(score, 0, "No match should score 0");
|
||||
}
|
||||
|
||||
/// Helper to create a registry with catalog entries (MCP servers come from catalog now).
|
||||
fn registry_with_catalog() -> ExtensionRegistry {
|
||||
let catalog = crate::registry::catalog::RegistryCatalog::load_or_embedded()
|
||||
.expect("catalog should load");
|
||||
let catalog_entries: Vec<RegistryEntry> = catalog
|
||||
.all()
|
||||
.iter()
|
||||
.filter_map(|m| m.to_registry_entry())
|
||||
.collect();
|
||||
ExtensionRegistry::new_with_catalog(catalog_entries)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_search_returns_sorted() {
|
||||
let registry = ExtensionRegistry::new();
|
||||
let registry = registry_with_catalog();
|
||||
let results = registry.search("notion").await;
|
||||
|
||||
assert!(!results.is_empty(), "Should find notion in registry");
|
||||
@@ -556,7 +381,7 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_search_empty_query_returns_all() {
|
||||
let registry = ExtensionRegistry::new();
|
||||
let registry = registry_with_catalog();
|
||||
let results = registry.search("").await;
|
||||
|
||||
assert!(results.len() > 5, "Empty query should return all entries");
|
||||
@@ -564,7 +389,7 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_search_by_keyword() {
|
||||
let registry = ExtensionRegistry::new();
|
||||
let registry = registry_with_catalog();
|
||||
let results = registry.search("issues tickets").await;
|
||||
|
||||
assert!(
|
||||
@@ -578,7 +403,7 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_exact_name() {
|
||||
let registry = ExtensionRegistry::new();
|
||||
let registry = registry_with_catalog();
|
||||
|
||||
let entry = registry.get("notion").await;
|
||||
assert!(entry.is_some());
|
||||
@@ -658,17 +483,30 @@ mod tests {
|
||||
auth_hint: AuthHint::CapabilitiesAuth,
|
||||
version: None,
|
||||
},
|
||||
// This shares a name with the builtin slack-mcp but has a different kind, so both should appear
|
||||
// Two entries with same name but different kinds should coexist
|
||||
RegistryEntry {
|
||||
name: "slack-mcp".to_string(),
|
||||
display_name: "Slack MCP WASM".to_string(),
|
||||
name: "dual-ext".to_string(),
|
||||
display_name: "Dual MCP".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description: "Dual extension MCP server".to_string(),
|
||||
keywords: vec!["messaging".into()],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://mcp.example.com".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "dual-ext".to_string(),
|
||||
display_name: "Dual WASM".to_string(),
|
||||
kind: ExtensionKind::WasmTool,
|
||||
description: "Slack WASM tool".to_string(),
|
||||
description: "Dual extension WASM tool".to_string(),
|
||||
keywords: vec!["messaging".into()],
|
||||
source: ExtensionSource::WasmBuildable {
|
||||
source_dir: "tools-src/slack".to_string(),
|
||||
build_dir: Some("tools-src/slack".to_string()),
|
||||
crate_name: Some("slack-tool".to_string()),
|
||||
source_dir: "tools-src/dual".to_string(),
|
||||
build_dir: Some("tools-src/dual".to_string()),
|
||||
crate_name: Some("dual-tool".to_string()),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::CapabilitiesAuth,
|
||||
@@ -683,41 +521,56 @@ mod tests {
|
||||
assert!(!results.is_empty(), "Should find telegram from catalog");
|
||||
assert_eq!(results[0].entry.name, "telegram");
|
||||
|
||||
// Should have both builtin MCP slack-mcp and catalog WASM slack-mcp
|
||||
let results = registry.search("slack").await;
|
||||
let slack_mcp = results
|
||||
// Should have both MCP and WASM entries with the same name
|
||||
let results = registry.search("dual-ext").await;
|
||||
let has_mcp = results
|
||||
.iter()
|
||||
.any(|r| r.entry.name == "slack-mcp" && r.entry.kind == ExtensionKind::McpServer);
|
||||
let slack_wasm = results
|
||||
.any(|r| r.entry.name == "dual-ext" && r.entry.kind == ExtensionKind::McpServer);
|
||||
let has_wasm = results
|
||||
.iter()
|
||||
.any(|r| r.entry.name == "slack-mcp" && r.entry.kind == ExtensionKind::WasmTool);
|
||||
assert!(slack_mcp, "Should have builtin MCP slack-mcp");
|
||||
assert!(slack_wasm, "Should have catalog WASM slack-mcp");
|
||||
.any(|r| r.entry.name == "dual-ext" && r.entry.kind == ExtensionKind::WasmTool);
|
||||
assert!(has_mcp, "Should have MCP dual-ext");
|
||||
assert!(has_wasm, "Should have WASM dual-ext");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_new_with_catalog_dedup_same_kind() {
|
||||
// A catalog entry with same name AND kind as a builtin should be skipped
|
||||
let catalog_entries = vec![RegistryEntry {
|
||||
name: "slack-mcp".to_string(),
|
||||
display_name: "Slack MCP Override".to_string(),
|
||||
kind: ExtensionKind::McpServer, // same kind as builtin slack-mcp
|
||||
description: "Should be skipped".to_string(),
|
||||
keywords: vec![],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://other.slack.com".to_string(),
|
||||
// When two catalog entries share name AND kind, only the first should be kept
|
||||
let catalog_entries = vec![
|
||||
RegistryEntry {
|
||||
name: "test-ext".to_string(),
|
||||
display_name: "Test First".to_string(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description: "First entry".to_string(),
|
||||
keywords: vec![],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://first.example.com".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
}];
|
||||
RegistryEntry {
|
||||
name: "test-ext".to_string(),
|
||||
display_name: "Test Duplicate".to_string(),
|
||||
kind: ExtensionKind::McpServer, // same kind
|
||||
description: "Should be skipped".to_string(),
|
||||
keywords: vec![],
|
||||
source: ExtensionSource::McpUrl {
|
||||
url: "https://second.example.com".to_string(),
|
||||
},
|
||||
fallback_source: None,
|
||||
auth_hint: AuthHint::Dcr,
|
||||
version: None,
|
||||
},
|
||||
];
|
||||
|
||||
let registry = ExtensionRegistry::new_with_catalog(catalog_entries);
|
||||
|
||||
let entry = registry.get("slack-mcp").await;
|
||||
let entry = registry.get("test-ext").await;
|
||||
assert!(entry.is_some());
|
||||
// Should still be the builtin, not the override
|
||||
assert_eq!(entry.unwrap().display_name, "Slack MCP");
|
||||
// Should be the first entry, not the duplicate
|
||||
assert_eq!(entry.unwrap().display_name, "Test First");
|
||||
}
|
||||
|
||||
#[tokio::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);
|
||||
}
|
||||
|
||||
|
||||
@@ -163,3 +163,42 @@ pub struct NearAiConfig {
|
||||
/// Enable cascade mode for smart routing. Default: true.
|
||||
pub smart_routing_cascade: bool,
|
||||
}
|
||||
|
||||
impl NearAiConfig {
|
||||
/// Create a minimal config suitable for listing available models.
|
||||
///
|
||||
/// Reads `NEARAI_API_KEY` from the environment and selects the
|
||||
/// appropriate base URL (cloud-api when API key is present,
|
||||
/// private.near.ai for session-token auth).
|
||||
pub(crate) fn for_model_discovery() -> Self {
|
||||
let api_key = std::env::var("NEARAI_API_KEY")
|
||||
.ok()
|
||||
.filter(|k| !k.is_empty())
|
||||
.map(SecretString::from);
|
||||
|
||||
let default_base = if api_key.is_some() {
|
||||
"https://cloud-api.near.ai"
|
||||
} else {
|
||||
"https://private.near.ai"
|
||||
};
|
||||
let base_url =
|
||||
std::env::var("NEARAI_BASE_URL").unwrap_or_else(|_| default_base.to_string());
|
||||
|
||||
Self {
|
||||
model: String::new(),
|
||||
cheap_model: None,
|
||||
base_url,
|
||||
api_key,
|
||||
fallback_model: None,
|
||||
max_retries: 3,
|
||||
circuit_breaker_threshold: None,
|
||||
circuit_breaker_recovery_secs: 30,
|
||||
response_cache_enabled: false,
|
||||
response_cache_ttl_secs: 3600,
|
||||
response_cache_max_entries: 1000,
|
||||
failover_cooldown_secs: 300,
|
||||
failover_cooldown_threshold: 3,
|
||||
smart_routing_cascade: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -29,6 +29,7 @@ pub mod session;
|
||||
pub mod smart_routing;
|
||||
|
||||
pub mod image_models;
|
||||
pub mod models;
|
||||
pub mod reasoning_models;
|
||||
pub mod vision_models;
|
||||
|
||||
|
||||
@@ -0,0 +1,349 @@
|
||||
//! Model discovery and fetching for multiple LLM providers.
|
||||
|
||||
/// Fetch models from the Anthropic API.
|
||||
///
|
||||
/// Returns `(model_id, display_label)` pairs. Falls back to static defaults on error.
|
||||
pub(crate) async fn fetch_anthropic_models(cached_key: Option<&str>) -> Vec<(String, String)> {
|
||||
let static_defaults = vec![
|
||||
(
|
||||
"claude-opus-4-6".into(),
|
||||
"Claude Opus 4.6 (latest flagship)".into(),
|
||||
),
|
||||
("claude-sonnet-4-6".into(), "Claude Sonnet 4.6".into()),
|
||||
("claude-opus-4-5".into(), "Claude Opus 4.5".into()),
|
||||
("claude-sonnet-4-5".into(), "Claude Sonnet 4.5".into()),
|
||||
("claude-haiku-4-5".into(), "Claude Haiku 4.5 (fast)".into()),
|
||||
];
|
||||
|
||||
let api_key = cached_key
|
||||
.map(String::from)
|
||||
.or_else(|| std::env::var("ANTHROPIC_API_KEY").ok())
|
||||
.filter(|k| !k.is_empty() && k != crate::config::OAUTH_PLACEHOLDER);
|
||||
|
||||
// Fall back to OAuth token if no API key
|
||||
let oauth_token = if api_key.is_none() {
|
||||
crate::config::helpers::optional_env("ANTHROPIC_OAUTH_TOKEN")
|
||||
.ok()
|
||||
.flatten()
|
||||
.filter(|t| !t.is_empty())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let (key_or_token, is_oauth) = match (api_key, oauth_token) {
|
||||
(Some(k), _) => (k, false),
|
||||
(None, Some(t)) => (t, true),
|
||||
(None, None) => return static_defaults,
|
||||
};
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let mut request = client
|
||||
.get("https://api.anthropic.com/v1/models")
|
||||
.header("anthropic-version", "2023-06-01")
|
||||
.timeout(std::time::Duration::from_secs(5));
|
||||
|
||||
if is_oauth {
|
||||
request = request
|
||||
.bearer_auth(&key_or_token)
|
||||
.header("anthropic-beta", "oauth-2025-04-20");
|
||||
} else {
|
||||
request = request.header("x-api-key", &key_or_token);
|
||||
}
|
||||
|
||||
let resp = match request.send().await {
|
||||
Ok(r) if r.status().is_success() => r,
|
||||
_ => return static_defaults,
|
||||
};
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct ModelEntry {
|
||||
id: String,
|
||||
}
|
||||
#[derive(serde::Deserialize)]
|
||||
struct ModelsResponse {
|
||||
data: Vec<ModelEntry>,
|
||||
}
|
||||
|
||||
match resp.json::<ModelsResponse>().await {
|
||||
Ok(body) => {
|
||||
let mut models: Vec<(String, String)> = body
|
||||
.data
|
||||
.into_iter()
|
||||
.filter(|m| !m.id.contains("embedding") && !m.id.contains("audio"))
|
||||
.map(|m| {
|
||||
let label = m.id.clone();
|
||||
(m.id, label)
|
||||
})
|
||||
.collect();
|
||||
if models.is_empty() {
|
||||
return static_defaults;
|
||||
}
|
||||
models.sort_by(|a, b| a.0.cmp(&b.0));
|
||||
models
|
||||
}
|
||||
Err(_) => static_defaults,
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch models from the OpenAI API.
|
||||
///
|
||||
/// Returns `(model_id, display_label)` pairs. Falls back to static defaults on error.
|
||||
pub(crate) async fn fetch_openai_models(cached_key: Option<&str>) -> Vec<(String, String)> {
|
||||
let static_defaults = vec![
|
||||
(
|
||||
"gpt-5.3-codex".into(),
|
||||
"GPT-5.3 Codex (latest flagship)".into(),
|
||||
),
|
||||
("gpt-5.2-codex".into(), "GPT-5.2 Codex".into()),
|
||||
("gpt-5.2".into(), "GPT-5.2".into()),
|
||||
(
|
||||
"gpt-5.1-codex-mini".into(),
|
||||
"GPT-5.1 Codex Mini (fast)".into(),
|
||||
),
|
||||
("gpt-5".into(), "GPT-5".into()),
|
||||
("gpt-5-mini".into(), "GPT-5 Mini".into()),
|
||||
("gpt-4.1".into(), "GPT-4.1".into()),
|
||||
("gpt-4.1-mini".into(), "GPT-4.1 Mini".into()),
|
||||
("o4-mini".into(), "o4-mini (fast reasoning)".into()),
|
||||
("o3".into(), "o3 (reasoning)".into()),
|
||||
];
|
||||
|
||||
let api_key = cached_key
|
||||
.map(String::from)
|
||||
.or_else(|| std::env::var("OPENAI_API_KEY").ok())
|
||||
.filter(|k| !k.is_empty());
|
||||
|
||||
let api_key = match api_key {
|
||||
Some(k) => k,
|
||||
None => return static_defaults,
|
||||
};
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let resp = match client
|
||||
.get("https://api.openai.com/v1/models")
|
||||
.bearer_auth(&api_key)
|
||||
.timeout(std::time::Duration::from_secs(5))
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(r) if r.status().is_success() => r,
|
||||
_ => return static_defaults,
|
||||
};
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct ModelEntry {
|
||||
id: String,
|
||||
}
|
||||
#[derive(serde::Deserialize)]
|
||||
struct ModelsResponse {
|
||||
data: Vec<ModelEntry>,
|
||||
}
|
||||
|
||||
match resp.json::<ModelsResponse>().await {
|
||||
Ok(body) => {
|
||||
let mut models: Vec<(String, String)> = body
|
||||
.data
|
||||
.into_iter()
|
||||
.filter(|m| is_openai_chat_model(&m.id))
|
||||
.map(|m| {
|
||||
let label = m.id.clone();
|
||||
(m.id, label)
|
||||
})
|
||||
.collect();
|
||||
if models.is_empty() {
|
||||
return static_defaults;
|
||||
}
|
||||
sort_openai_models(&mut models);
|
||||
models
|
||||
}
|
||||
Err(_) => static_defaults,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_openai_chat_model(model_id: &str) -> bool {
|
||||
let id = model_id.to_ascii_lowercase();
|
||||
|
||||
let is_chat_family = id.starts_with("gpt-")
|
||||
|| id.starts_with("chatgpt-")
|
||||
|| id.starts_with("o1")
|
||||
|| id.starts_with("o3")
|
||||
|| id.starts_with("o4")
|
||||
|| id.starts_with("o5");
|
||||
|
||||
let is_non_chat_variant = id.contains("realtime")
|
||||
|| id.contains("audio")
|
||||
|| id.contains("transcribe")
|
||||
|| id.contains("tts")
|
||||
|| id.contains("embedding")
|
||||
|| id.contains("moderation")
|
||||
|| id.contains("image");
|
||||
|
||||
is_chat_family && !is_non_chat_variant
|
||||
}
|
||||
|
||||
pub(crate) fn openai_model_priority(model_id: &str) -> usize {
|
||||
let id = model_id.to_ascii_lowercase();
|
||||
|
||||
const EXACT_PRIORITY: &[&str] = &[
|
||||
"gpt-5.3-codex",
|
||||
"gpt-5.2-codex",
|
||||
"gpt-5.2",
|
||||
"gpt-5.1-codex-mini",
|
||||
"gpt-5",
|
||||
"gpt-5-mini",
|
||||
"gpt-5-nano",
|
||||
"o4-mini",
|
||||
"o3",
|
||||
"o1",
|
||||
"gpt-4.1",
|
||||
"gpt-4.1-mini",
|
||||
"gpt-4o",
|
||||
"gpt-4o-mini",
|
||||
];
|
||||
if let Some(pos) = EXACT_PRIORITY.iter().position(|m| id == *m) {
|
||||
return pos;
|
||||
}
|
||||
|
||||
const PREFIX_PRIORITY: &[&str] = &[
|
||||
"gpt-5.", "gpt-5-", "o3-", "o4-", "o1-", "gpt-4.1-", "gpt-4o-", "gpt-3.5-", "chatgpt-",
|
||||
];
|
||||
if let Some(pos) = PREFIX_PRIORITY
|
||||
.iter()
|
||||
.position(|prefix| id.starts_with(prefix))
|
||||
{
|
||||
return EXACT_PRIORITY.len() + pos;
|
||||
}
|
||||
|
||||
EXACT_PRIORITY.len() + PREFIX_PRIORITY.len() + 1
|
||||
}
|
||||
|
||||
pub(crate) fn sort_openai_models(models: &mut [(String, String)]) {
|
||||
models.sort_by(|a, b| {
|
||||
openai_model_priority(&a.0)
|
||||
.cmp(&openai_model_priority(&b.0))
|
||||
.then_with(|| a.0.cmp(&b.0))
|
||||
});
|
||||
}
|
||||
|
||||
/// Fetch installed models from a local Ollama instance.
|
||||
///
|
||||
/// Returns `(model_name, display_label)` pairs. Falls back to static defaults on error.
|
||||
pub(crate) async fn fetch_ollama_models(base_url: &str) -> Vec<(String, String)> {
|
||||
let static_defaults = vec![
|
||||
("llama3".into(), "llama3".into()),
|
||||
("mistral".into(), "mistral".into()),
|
||||
("codellama".into(), "codellama".into()),
|
||||
];
|
||||
|
||||
let url = format!("{}/api/tags", base_url.trim_end_matches('/'));
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let resp = match client
|
||||
.get(&url)
|
||||
.timeout(std::time::Duration::from_secs(5))
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(r) if r.status().is_success() => r,
|
||||
Ok(_) => return static_defaults,
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
"Could not connect to Ollama at {base_url}. Is it running? Using static defaults."
|
||||
);
|
||||
return static_defaults;
|
||||
}
|
||||
};
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct ModelEntry {
|
||||
name: String,
|
||||
}
|
||||
#[derive(serde::Deserialize)]
|
||||
struct TagsResponse {
|
||||
models: Vec<ModelEntry>,
|
||||
}
|
||||
|
||||
match resp.json::<TagsResponse>().await {
|
||||
Ok(body) => {
|
||||
let models: Vec<(String, String)> = body
|
||||
.models
|
||||
.into_iter()
|
||||
.map(|m| {
|
||||
let label = m.name.clone();
|
||||
(m.name, label)
|
||||
})
|
||||
.collect();
|
||||
if models.is_empty() {
|
||||
return static_defaults;
|
||||
}
|
||||
models
|
||||
}
|
||||
Err(_) => static_defaults,
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch models from a generic OpenAI-compatible /v1/models endpoint.
|
||||
///
|
||||
/// Used for registry providers like Groq, NVIDIA NIM, etc.
|
||||
pub(crate) async fn fetch_openai_compatible_models(
|
||||
base_url: &str,
|
||||
cached_key: Option<&str>,
|
||||
) -> Vec<(String, String)> {
|
||||
if base_url.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
|
||||
let url = format!("{}/models", base_url.trim_end_matches('/'));
|
||||
let client = reqwest::Client::new();
|
||||
let mut req = client.get(&url).timeout(std::time::Duration::from_secs(5));
|
||||
if let Some(key) = cached_key {
|
||||
req = req.bearer_auth(key);
|
||||
}
|
||||
|
||||
let resp = match req.send().await {
|
||||
Ok(r) if r.status().is_success() => r,
|
||||
_ => return vec![],
|
||||
};
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Model {
|
||||
id: String,
|
||||
}
|
||||
#[derive(serde::Deserialize)]
|
||||
struct ModelsResponse {
|
||||
data: Vec<Model>,
|
||||
}
|
||||
|
||||
match resp.json::<ModelsResponse>().await {
|
||||
Ok(body) => body
|
||||
.data
|
||||
.into_iter()
|
||||
.map(|m| {
|
||||
let label = m.id.clone();
|
||||
(m.id, label)
|
||||
})
|
||||
.collect(),
|
||||
Err(_) => vec![],
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the `LlmConfig` used by `fetch_nearai_models` to list available models.
|
||||
///
|
||||
/// Uses [`NearAiConfig::for_model_discovery()`] to construct a minimal NEAR AI
|
||||
/// config, then wraps it in an `LlmConfig` with session config for auth.
|
||||
pub(crate) fn build_nearai_model_fetch_config() -> crate::config::LlmConfig {
|
||||
let auth_base_url =
|
||||
std::env::var("NEARAI_AUTH_URL").unwrap_or_else(|_| "https://private.near.ai".to_string());
|
||||
|
||||
crate::config::LlmConfig {
|
||||
backend: "nearai".to_string(),
|
||||
session: crate::llm::session::SessionConfig {
|
||||
auth_base_url,
|
||||
session_path: crate::config::llm::default_session_path(),
|
||||
},
|
||||
nearai: crate::config::NearAiConfig::for_model_discovery(),
|
||||
provider: None,
|
||||
bedrock: None,
|
||||
request_timeout_secs: 120,
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -155,22 +155,22 @@ pub fn is_silent_reply(text: &str) -> bool {
|
||||
|
||||
/// Quick-check: bail early if no reasoning/final tags are present at all.
|
||||
static QUICK_TAG_RE: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r"(?i)<\s*/?\s*(?:think(?:ing)?|thought|thoughts|antthinking|reasoning|reflection|scratchpad|inner_monologue|final)\b").expect("QUICK_TAG_RE")
|
||||
Regex::new(r"(?i)<\s*/?\s*(?:think(?:ing)?|thought|thoughts|antthinking|reasoning|reflection|scratchpad|inner_monologue|final)\b").expect("QUICK_TAG_RE") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
/// Matches thinking/reasoning open and close tags. Capture group 1 is "/" for close tags.
|
||||
/// Whitespace-tolerant, case-insensitive, attribute-aware.
|
||||
static THINKING_TAG_RE: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r"(?i)<\s*(/?)\s*(?:think(?:ing)?|thought|thoughts|antthinking|reasoning|reflection|scratchpad|inner_monologue)\b[^<>]*>").expect("THINKING_TAG_RE")
|
||||
Regex::new(r"(?i)<\s*(/?)\s*(?:think(?:ing)?|thought|thoughts|antthinking|reasoning|reflection|scratchpad|inner_monologue)\b[^<>]*>").expect("THINKING_TAG_RE") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
/// Matches `<final>` / `</final>` tags. Capture group 1 is "/" for close tags.
|
||||
static FINAL_TAG_RE: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"(?i)<\s*(/?)\s*final\b[^<>]*>").expect("FINAL_TAG_RE"));
|
||||
LazyLock::new(|| Regex::new(r"(?i)<\s*(/?)\s*final\b[^<>]*>").expect("FINAL_TAG_RE")); // safety: hardcoded literal
|
||||
|
||||
/// Matches pipe-delimited reasoning tags: `<|think|>...<|/think|>` etc.
|
||||
static PIPE_REASONING_TAG_RE: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r"(?i)<\|(/?)\s*(?:think(?:ing)?|thought|thoughts|antthinking|reasoning|reflection|scratchpad|inner_monologue)\|>").expect("PIPE_REASONING_TAG_RE")
|
||||
Regex::new(r"(?i)<\|(/?)\s*(?:think(?:ing)?|thought|thoughts|antthinking|reasoning|reflection|scratchpad|inner_monologue)\|>").expect("PIPE_REASONING_TAG_RE") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
/// Context for reasoning operations.
|
||||
|
||||
+1
-1
@@ -219,7 +219,7 @@ impl ProviderRegistry {
|
||||
pub fn load() -> Self {
|
||||
let builtins: Vec<ProviderDefinition> =
|
||||
serde_json::from_str(include_str!("../../providers.json"))
|
||||
.expect("built-in providers.json must be valid JSON");
|
||||
.expect("built-in providers.json must be valid JSON"); // safety: compile-time embedded file
|
||||
|
||||
let mut all = builtins;
|
||||
|
||||
|
||||
@@ -548,6 +548,7 @@ mod tests {
|
||||
model: None,
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
stop_sequences: None,
|
||||
tool_choice: None,
|
||||
metadata: Default::default(),
|
||||
};
|
||||
|
||||
+22
-21
@@ -248,7 +248,7 @@ fn build_domain_regex(keywords: &[&str]) -> Regex {
|
||||
let pattern = format!(r"(?i)\b({})\b", keywords.join("|"));
|
||||
Regex::new(&pattern).unwrap_or_else(|e| {
|
||||
tracing::warn!(error = %e, "Invalid domain keywords pattern, using minimal fallback");
|
||||
Regex::new(r"(?i)\b(api|code|deploy)\b").expect("fallback regex is valid")
|
||||
Regex::new(r"(?i)\b(api|code|deploy)\b").expect("fallback regex is valid") // safety: hardcoded literal
|
||||
})
|
||||
}
|
||||
|
||||
@@ -274,71 +274,71 @@ use std::sync::LazyLock;
|
||||
static RE_REASONING: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(
|
||||
r"(?i)\b(why|how|explain|analyze|analyse|compare|contrast|evaluate|assess|reason|think|consider|implications?|consequences?|trade-?offs?|pros?\s*(and|&)\s*cons?|advantages?|disadvantages?|benefits?|drawbacks?|differs?|difference|versus|vs\.?|better|worse|optimal|best|worst)\b"
|
||||
).expect("RE_REASONING is a valid regex")
|
||||
).expect("RE_REASONING is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_MULTI_STEP: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(
|
||||
r"(?i)\b(first|then|next|after|before|finally|step|steps|phase|stages?|process|workflow|sequence|procedure|pipeline|chain|series|order|followed by)\b"
|
||||
).expect("RE_MULTI_STEP is a valid regex")
|
||||
).expect("RE_MULTI_STEP is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_CREATIVITY: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(
|
||||
r"(?i)\b(write|create|generate|compose|design|imagine|brainstorm|ideate|draft|invent|story|poem|essay|article|blog|content|narrative|script|summarize|summarise|rewrite|paraphrase|translate|adapt|tweet|post|thread|outline|structure|format|style|tone|voice)\b"
|
||||
).expect("RE_CREATIVITY is a valid regex")
|
||||
).expect("RE_CREATIVITY is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_PRECISION: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(
|
||||
r"(?i)\b(\d{4}|\d+\.\d+|exactly|precisely|specific|accurate|correct|verify|confirm|date|time|number|calculate|compute|measure|count)\b"
|
||||
).expect("RE_PRECISION is a valid regex")
|
||||
).expect("RE_PRECISION is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_CODE: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(
|
||||
r"(?i)(`{1,3}|```|function|const|let|var|import|export|class|def |async|await|=>|\.ts|\.js|\.py|\.rs|\.go|\.sol|\(\)|\[\]|\{\}|<[A-Z][a-z]+>|useState|useEffect|npm|yarn|pnpm|cargo|pip|implement|rebase|merge|commit|branch|PR|pull.?request|columns?|migrations?|module|refactor|debug|fix|bug|error|schema|database|query)"
|
||||
).expect("RE_CODE is a valid regex")
|
||||
).expect("RE_CODE is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_TOOL: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(
|
||||
r"(?i)\b(file|read|write|search|fetch|run|execute|check|look up|find|open|save|send|post|get|download|upload|install|deploy|build|compile|test|add|update|remove|delete|modify|change|edit|create|resolve|push|pull|clone)\b"
|
||||
).expect("RE_TOOL is a valid regex")
|
||||
).expect("RE_TOOL is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_SAFETY: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(
|
||||
r"(?i)\b(password|secret|private|confidential|medical|legal|financial|personal|sensitive|ssn|credit.?card|auth|token|key|encrypt|decrypt|hash|vulnerability|exploit|attack|breach)\b"
|
||||
).expect("RE_SAFETY is a valid regex")
|
||||
).expect("RE_SAFETY is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_CONTEXT: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(
|
||||
r"(?i)\b(previous|earlier|above|before|last|that|those|it|they|we discussed|you said|mentioned|remember|recall|as I said|like I mentioned)\b"
|
||||
).expect("RE_CONTEXT is a valid regex")
|
||||
).expect("RE_CONTEXT is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_VAGUE: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r"(?i)\b(it|this|that|something|stuff|thing|things)\b")
|
||||
.expect("RE_VAGUE is a valid regex")
|
||||
.expect("RE_VAGUE is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_OPEN_ENDED: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r"(?i)\b(why|how|what if|explain|describe|elaborate|discuss)\b")
|
||||
.expect("RE_OPEN_ENDED is a valid regex")
|
||||
.expect("RE_OPEN_ENDED is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_CONJUNCTIONS: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(
|
||||
r"(?i)\b(and|but|or|however|therefore|because|although|while|whereas|moreover|furthermore)\b",
|
||||
)
|
||||
.expect("RE_CONJUNCTIONS is a valid regex")
|
||||
.expect("RE_CONJUNCTIONS is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
static RE_TIER_HINT: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r"(?i)\[tier:(flash|standard|pro|frontier)\]")
|
||||
.expect("RE_TIER_HINT is a valid regex")
|
||||
.expect("RE_TIER_HINT is a valid regex") // safety: hardcoded literal
|
||||
});
|
||||
|
||||
/// Default domain regex, compiled once from `DEFAULT_DOMAIN_KEYWORDS`.
|
||||
@@ -363,7 +363,7 @@ static DEFAULT_OVERRIDES: LazyLock<Vec<PatternOverride>> = LazyLock::new(|| {
|
||||
regex: Regex::new(
|
||||
r"(?i)^(hi|hello|hey|thanks|ok|sure|yes|no|yep|nope|cool|nice|great|got it)$",
|
||||
)
|
||||
.expect("greeting pattern is valid"),
|
||||
.expect("greeting pattern is valid"), // safety: hardcoded literal
|
||||
tier: Tier::Flash,
|
||||
},
|
||||
// Flash tier: quick lookups (end-anchored to avoid matching complex questions
|
||||
@@ -372,29 +372,29 @@ static DEFAULT_OVERRIDES: LazyLock<Vec<PatternOverride>> = LazyLock::new(|| {
|
||||
regex: Regex::new(
|
||||
r"(?i)^what(?:'s|\s+is)?\s+(?:the\s+)?(time|date|day|weather)\b(?:\s+(?:is\s+it|today|now|in\s+\S+))?[?.!]*$",
|
||||
)
|
||||
.expect("lookup pattern is valid"),
|
||||
.expect("lookup pattern is valid"), // safety: hardcoded literal
|
||||
tier: Tier::Flash,
|
||||
},
|
||||
// Frontier tier: security audits
|
||||
PatternOverride {
|
||||
regex: Regex::new(r"(?i)security.*(audit|review|scan)")
|
||||
.expect("security audit pattern is valid"),
|
||||
.expect("security audit pattern is valid"), // safety: hardcoded literal
|
||||
tier: Tier::Frontier,
|
||||
},
|
||||
PatternOverride {
|
||||
regex: Regex::new(r"(?i)vulnerabilit(y|ies).*(review|scan|check|audit)")
|
||||
.expect("vulnerability pattern is valid"),
|
||||
.expect("vulnerability pattern is valid"), // safety: hardcoded literal
|
||||
tier: Tier::Frontier,
|
||||
},
|
||||
// Pro tier: production deployments
|
||||
PatternOverride {
|
||||
regex: Regex::new(r"(?i)deploy.*(mainnet|production)")
|
||||
.expect("deploy pattern is valid"),
|
||||
.expect("deploy pattern is valid"), // safety: hardcoded literal
|
||||
tier: Tier::Pro,
|
||||
},
|
||||
PatternOverride {
|
||||
regex: Regex::new(r"(?i)production.*(deploy|release|push)")
|
||||
.expect("production pattern is valid"),
|
||||
.expect("production pattern is valid"), // safety: hardcoded literal
|
||||
tier: Tier::Pro,
|
||||
},
|
||||
]
|
||||
@@ -451,7 +451,7 @@ fn score_complexity_internal(
|
||||
|
||||
// Check for explicit tier hint (e.g. "[tier:flash]")
|
||||
if let Some(caps) = RE_TIER_HINT.captures(prompt) {
|
||||
let tier_str = caps.get(1).expect("capture group 1 exists").as_str();
|
||||
let tier_str = caps.get(1).expect("capture group 1 exists").as_str(); // safety: RE_TIER_HINT has group 1
|
||||
let tier = match tier_str.to_lowercase().as_str() {
|
||||
"flash" => Tier::Flash,
|
||||
"standard" => Tier::Standard,
|
||||
@@ -758,7 +758,8 @@ impl SmartRoutingProvider {
|
||||
|
||||
// Highest priority: explicit tier hints (e.g. "[tier:flash]")
|
||||
if let Some(caps) = RE_TIER_HINT.captures(last_user_msg) {
|
||||
let tier_str = caps.get(1).expect("capture group 1 exists").as_str();
|
||||
// SAFETY: RE_TIER_HINT has exactly one capture group; get(1) is guaranteed Some after match.
|
||||
let tier_str = caps.get(1).expect("capture group 1 exists").as_str(); // safety: RE_TIER_HINT has group 1
|
||||
let tier = match tier_str.to_lowercase().as_str() {
|
||||
"flash" => Tier::Flash,
|
||||
"standard" => Tier::Standard,
|
||||
|
||||
+14
-1
@@ -92,6 +92,10 @@ async fn async_main() -> anyhow::Result<()> {
|
||||
return ironclaw::cli::run_skills_command(skills_cmd.clone(), cli.config.as_deref())
|
||||
.await;
|
||||
}
|
||||
Some(Command::Logs(logs_cmd)) => {
|
||||
init_cli_tracing();
|
||||
return ironclaw::cli::run_logs_command(logs_cmd.clone(), cli.config.as_deref()).await;
|
||||
}
|
||||
Some(Command::Doctor) => {
|
||||
init_cli_tracing();
|
||||
return ironclaw::cli::run_doctor_command().await;
|
||||
@@ -920,7 +924,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 {
|
||||
|
||||
@@ -528,6 +528,169 @@ 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 super::*;
|
||||
@@ -820,4 +983,38 @@ mod tests {
|
||||
_ => panic!("expected Lightweight"),
|
||||
}
|
||||
}
|
||||
|
||||
#[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
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_describe_cron_edge_cases() {
|
||||
assert_eq!(describe_cron("", None), "cron: (empty)"); // safety: test-only
|
||||
assert_eq!(describe_cron("not a cron", None), "cron: not a cron"); // safety: test-only
|
||||
let weekdays_5_field = describe_cron("0 9 * * MON-FRI", None);
|
||||
assert_eq!(weekdays_5_field, "Weekdays at 9:00 AM"); // safety: test-only
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
};
|
||||
|
||||
+86
-31
@@ -192,6 +192,12 @@ impl RegistryCatalog {
|
||||
Self::load_manifests_from_dir(&channels_dir, "channels", &mut manifests)?;
|
||||
}
|
||||
|
||||
// Load MCP servers
|
||||
let mcp_servers_dir = registry_dir.join("mcp-servers");
|
||||
if mcp_servers_dir.is_dir() {
|
||||
Self::load_manifests_from_dir(&mcp_servers_dir, "mcp-servers", &mut manifests)?;
|
||||
}
|
||||
|
||||
// Load bundles
|
||||
let bundles_path = registry_dir.join("_bundles.json");
|
||||
let bundles = if bundles_path.is_file() {
|
||||
@@ -280,8 +286,9 @@ impl RegistryCatalog {
|
||||
/// Get a manifest by name. Tries exact key match first ("tools/github"),
|
||||
/// then searches by bare name ("github").
|
||||
///
|
||||
/// If a bare name matches both a tool and a channel, returns `None`.
|
||||
/// Use a qualified key ("tools/github" or "channels/telegram") to disambiguate.
|
||||
/// If a bare name matches more than one prefix, returns `None`.
|
||||
/// Use a qualified key ("tools/github", "channels/telegram", or
|
||||
/// "mcp-servers/notion") to disambiguate.
|
||||
pub fn get(&self, name: &str) -> Option<&ExtensionManifest> {
|
||||
// Try exact key first
|
||||
if let Some(m) = self.manifests.get(name) {
|
||||
@@ -289,14 +296,15 @@ impl RegistryCatalog {
|
||||
}
|
||||
|
||||
// Try with kind prefix, detecting collisions
|
||||
let tool = self.manifests.get(&format!("tools/{}", name));
|
||||
let channel = self.manifests.get(&format!("channels/{}", name));
|
||||
let candidates: Vec<_> = ["tools", "channels", "mcp-servers"]
|
||||
.iter()
|
||||
.filter_map(|prefix| self.manifests.get(&format!("{}/{}", prefix, name)))
|
||||
.collect();
|
||||
|
||||
match (tool, channel) {
|
||||
(Some(_), Some(_)) => None, // ambiguous
|
||||
(Some(m), None) => Some(m),
|
||||
(None, Some(m)) => Some(m),
|
||||
(None, None) => None,
|
||||
if candidates.len() == 1 {
|
||||
Some(candidates[0])
|
||||
} else {
|
||||
None // ambiguous or not found
|
||||
}
|
||||
}
|
||||
|
||||
@@ -308,37 +316,63 @@ impl RegistryCatalog {
|
||||
return Ok(m);
|
||||
}
|
||||
|
||||
let has_tool = self.manifests.contains_key(&format!("tools/{}", name));
|
||||
let has_channel = self.manifests.contains_key(&format!("channels/{}", name));
|
||||
let prefixes: &[(&str, &str)] = &[
|
||||
("tools", "tool"),
|
||||
("channels", "channel"),
|
||||
("mcp-servers", "mcp_server"),
|
||||
];
|
||||
|
||||
match (has_tool, has_channel) {
|
||||
(true, true) => Err(RegistryError::AmbiguousName {
|
||||
name: name.to_string(),
|
||||
kind_a: "tool",
|
||||
prefix_a: "tools",
|
||||
kind_b: "channel",
|
||||
prefix_b: "channels",
|
||||
}),
|
||||
(true, false) => Ok(self.manifests.get(&format!("tools/{}", name)).unwrap()),
|
||||
(false, true) => Ok(self.manifests.get(&format!("channels/{}", name)).unwrap()),
|
||||
(false, false) => Err(RegistryError::ExtensionNotFound(name.to_string())),
|
||||
let matches: Vec<_> = prefixes
|
||||
.iter()
|
||||
.filter(|(prefix, _)| self.manifests.contains_key(&format!("{}/{}", prefix, name)))
|
||||
.collect();
|
||||
|
||||
match matches.len() {
|
||||
0 => Err(RegistryError::ExtensionNotFound(name.to_string())),
|
||||
1 => {
|
||||
let (prefix, _) = matches[0];
|
||||
let key = format!("{}/{}", prefix, name);
|
||||
self.manifests
|
||||
.get(&key)
|
||||
.ok_or_else(|| RegistryError::ExtensionNotFound(name.to_string()))
|
||||
}
|
||||
_ => {
|
||||
let (prefix_a, kind_a) = matches[0];
|
||||
let (prefix_b, kind_b) = matches[1];
|
||||
Err(RegistryError::AmbiguousName {
|
||||
name: name.to_string(),
|
||||
kind_a,
|
||||
prefix_a,
|
||||
kind_b,
|
||||
prefix_b,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the full key ("tools/github" or "channels/telegram") for a manifest.
|
||||
/// Get the full key ("tools/github", "channels/telegram", or
|
||||
/// "mcp-servers/notion") for a manifest.
|
||||
pub fn key_for(&self, name: &str) -> Option<String> {
|
||||
if self.manifests.contains_key(name) {
|
||||
return Some(name.to_string());
|
||||
}
|
||||
|
||||
let has_tool = self.manifests.contains_key(&format!("tools/{}", name));
|
||||
let has_channel = self.manifests.contains_key(&format!("channels/{}", name));
|
||||
let matches: Vec<String> = ["tools", "channels", "mcp-servers"]
|
||||
.iter()
|
||||
.filter_map(|prefix| {
|
||||
let key = format!("{}/{}", prefix, name);
|
||||
if self.manifests.contains_key(&key) {
|
||||
Some(key)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
match (has_tool, has_channel) {
|
||||
(true, true) => None, // ambiguous
|
||||
(true, false) => Some(format!("tools/{}", name)),
|
||||
(false, true) => Some(format!("channels/{}", name)),
|
||||
(false, false) => None,
|
||||
if matches.len() == 1 {
|
||||
matches.into_iter().next()
|
||||
} else {
|
||||
None // ambiguous or not found
|
||||
}
|
||||
}
|
||||
|
||||
@@ -476,8 +510,10 @@ mod tests {
|
||||
fn create_test_registry(dir: &Path) {
|
||||
let tools_dir = dir.join("tools");
|
||||
let channels_dir = dir.join("channels");
|
||||
let mcp_dir = dir.join("mcp-servers");
|
||||
fs::create_dir_all(&tools_dir).unwrap();
|
||||
fs::create_dir_all(&channels_dir).unwrap();
|
||||
fs::create_dir_all(&mcp_dir).unwrap();
|
||||
|
||||
fs::write(
|
||||
tools_dir.join("slack.json"),
|
||||
@@ -540,6 +576,20 @@ mod tests {
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
fs::write(
|
||||
mcp_dir.join("notion.json"),
|
||||
r#"{
|
||||
"name": "notion",
|
||||
"display_name": "Notion",
|
||||
"kind": "mcp_server",
|
||||
"description": "Connect to Notion for pages and databases",
|
||||
"keywords": ["notes", "wiki"],
|
||||
"url": "https://mcp.notion.com/mcp",
|
||||
"auth": "dcr"
|
||||
}"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
fs::write(
|
||||
dir.join("_bundles.json"),
|
||||
r#"{
|
||||
@@ -565,7 +615,7 @@ mod tests {
|
||||
create_test_registry(tmp.path());
|
||||
|
||||
let catalog = RegistryCatalog::load(tmp.path()).unwrap();
|
||||
assert_eq!(catalog.all().len(), 3);
|
||||
assert_eq!(catalog.all().len(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -579,6 +629,9 @@ mod tests {
|
||||
|
||||
let channels = catalog.list(Some(ManifestKind::Channel), None);
|
||||
assert_eq!(channels.len(), 1);
|
||||
|
||||
let mcp_servers = catalog.list(Some(ManifestKind::McpServer), None);
|
||||
assert_eq!(mcp_servers.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -603,10 +656,12 @@ mod tests {
|
||||
|
||||
// Full key
|
||||
assert!(catalog.get("tools/slack").is_some());
|
||||
assert!(catalog.get("mcp-servers/notion").is_some());
|
||||
|
||||
// Bare name
|
||||
assert!(catalog.get("slack").is_some());
|
||||
assert!(catalog.get("telegram").is_some());
|
||||
assert!(catalog.get("notion").is_some());
|
||||
|
||||
// Missing
|
||||
assert!(catalog.get("nonexistent").is_none());
|
||||
|
||||
@@ -20,6 +20,8 @@ struct EmbeddedCatalogRaw {
|
||||
#[serde(default)]
|
||||
channels: Vec<ExtensionManifest>,
|
||||
#[serde(default)]
|
||||
mcp_servers: Vec<ExtensionManifest>,
|
||||
#[serde(default)]
|
||||
bundles: BundlesFile,
|
||||
}
|
||||
|
||||
@@ -52,6 +54,10 @@ fn parsed_catalog() -> &'static ParsedCatalog {
|
||||
let key = format!("channels/{}", m.name);
|
||||
manifests.insert(key, m);
|
||||
}
|
||||
for m in raw.mcp_servers {
|
||||
let key = format!("mcp-servers/{}", m.name);
|
||||
manifests.insert(key, m);
|
||||
}
|
||||
|
||||
ParsedCatalog {
|
||||
manifests,
|
||||
|
||||
+76
-18
@@ -7,7 +7,7 @@ use tokio::fs;
|
||||
|
||||
use crate::bootstrap::ironclaw_base_dir;
|
||||
use crate::registry::catalog::RegistryError;
|
||||
use crate::registry::manifest::{BundleDefinition, ExtensionManifest, ManifestKind};
|
||||
use crate::registry::manifest::{BundleDefinition, ExtensionManifest, ManifestKind, SourceSpec};
|
||||
|
||||
// GitHub-only by design. New trusted hosts (e.g. a NEAR AI CDN) must be
|
||||
// explicitly added here; unknown hosts fall back to source build with a
|
||||
@@ -98,12 +98,29 @@ fn validate_manifest_install_inputs(manifest: &ExtensionManifest) -> Result<(),
|
||||
});
|
||||
}
|
||||
|
||||
// MCP servers are not installed via this path
|
||||
if manifest.kind == ManifestKind::McpServer {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let source = match &manifest.source {
|
||||
Some(s) => s,
|
||||
None => {
|
||||
return Err(RegistryError::InvalidManifest {
|
||||
name: manifest.name.clone(),
|
||||
field: "source",
|
||||
reason: "WASM extensions must have a source spec".to_string(),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
let expected_prefix = match manifest.kind {
|
||||
ManifestKind::Tool => "tools-src/",
|
||||
ManifestKind::Channel => "channels-src/",
|
||||
ManifestKind::McpServer => unreachable!(),
|
||||
};
|
||||
|
||||
if !manifest.source.dir.starts_with(expected_prefix) {
|
||||
if !source.dir.starts_with(expected_prefix) {
|
||||
return Err(RegistryError::InvalidManifest {
|
||||
name: manifest.name.clone(),
|
||||
field: "source.dir",
|
||||
@@ -111,7 +128,7 @@ fn validate_manifest_install_inputs(manifest: &ExtensionManifest) -> Result<(),
|
||||
});
|
||||
}
|
||||
|
||||
let source_path = Path::new(&manifest.source.dir);
|
||||
let source_path = Path::new(&source.dir);
|
||||
let has_unsafe_component = source_path.components().any(|component| {
|
||||
matches!(
|
||||
component,
|
||||
@@ -127,9 +144,9 @@ fn validate_manifest_install_inputs(manifest: &ExtensionManifest) -> Result<(),
|
||||
});
|
||||
}
|
||||
|
||||
let has_path_separator = manifest.source.capabilities.contains('/')
|
||||
|| manifest.source.capabilities.contains('\\')
|
||||
|| manifest.source.capabilities.contains("..");
|
||||
let has_path_separator = source.capabilities.contains('/')
|
||||
|| source.capabilities.contains('\\')
|
||||
|| source.capabilities.contains("..");
|
||||
|
||||
if has_path_separator {
|
||||
return Err(RegistryError::InvalidManifest {
|
||||
@@ -142,6 +159,18 @@ fn validate_manifest_install_inputs(manifest: &ExtensionManifest) -> Result<(),
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Extract the source spec from a manifest, returning an error if absent.
|
||||
fn require_source(manifest: &ExtensionManifest) -> Result<&SourceSpec, RegistryError> {
|
||||
manifest
|
||||
.source
|
||||
.as_ref()
|
||||
.ok_or_else(|| RegistryError::InvalidManifest {
|
||||
name: manifest.name.clone(),
|
||||
field: "source",
|
||||
reason: "WASM extensions must have a source spec".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn download_failure_reason(error: &reqwest::Error) -> String {
|
||||
if error.is_timeout() {
|
||||
"request timed out".to_string()
|
||||
@@ -206,7 +235,17 @@ impl RegistryInstaller {
|
||||
) -> Result<InstallOutcome, RegistryError> {
|
||||
validate_manifest_install_inputs(manifest)?;
|
||||
|
||||
let source_dir = self.repo_root.join(&manifest.source.dir);
|
||||
if manifest.kind == ManifestKind::McpServer {
|
||||
return Err(RegistryError::InvalidManifest {
|
||||
name: manifest.name.clone(),
|
||||
field: "kind",
|
||||
reason: "MCP servers cannot be installed from source".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
let source = require_source(manifest)?;
|
||||
|
||||
let source_dir = self.repo_root.join(&source.dir);
|
||||
if !source_dir.exists() {
|
||||
return Err(RegistryError::ManifestRead {
|
||||
path: source_dir.clone(),
|
||||
@@ -217,6 +256,7 @@ impl RegistryInstaller {
|
||||
let target_dir = match manifest.kind {
|
||||
ManifestKind::Tool => &self.tools_dir,
|
||||
ManifestKind::Channel => &self.channels_dir,
|
||||
ManifestKind::McpServer => unreachable!(),
|
||||
};
|
||||
|
||||
fs::create_dir_all(target_dir)
|
||||
@@ -242,7 +282,7 @@ impl RegistryInstaller {
|
||||
manifest.display_name,
|
||||
source_dir.display()
|
||||
);
|
||||
let crate_name = &manifest.source.crate_name;
|
||||
let crate_name = &source.crate_name;
|
||||
let wasm_path =
|
||||
crate::registry::artifacts::build_wasm_component(&source_dir, crate_name, true)
|
||||
.await
|
||||
@@ -258,7 +298,7 @@ impl RegistryInstaller {
|
||||
.map_err(RegistryError::Io)?;
|
||||
|
||||
// Copy capabilities file
|
||||
let caps_source = source_dir.join(&manifest.source.capabilities);
|
||||
let caps_source = source_dir.join(&source.capabilities);
|
||||
let target_caps = target_dir.join(format!("{}.capabilities.json", manifest.name));
|
||||
let has_capabilities = if caps_source.exists() {
|
||||
fs::copy(&caps_source, &target_caps)
|
||||
@@ -296,6 +336,16 @@ impl RegistryInstaller {
|
||||
// catch it first.
|
||||
validate_manifest_install_inputs(manifest)?;
|
||||
|
||||
if manifest.kind == ManifestKind::McpServer {
|
||||
return Err(RegistryError::InvalidManifest {
|
||||
name: manifest.name.clone(),
|
||||
field: "kind",
|
||||
reason: "MCP servers cannot be installed via the WASM installer".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
let source = require_source(manifest)?;
|
||||
|
||||
let has_artifact = manifest
|
||||
.artifacts
|
||||
.get("wasm32-wasip2")
|
||||
@@ -306,7 +356,7 @@ impl RegistryInstaller {
|
||||
return self.install_from_source(manifest, force).await;
|
||||
}
|
||||
|
||||
let source_dir = self.repo_root.join(&manifest.source.dir);
|
||||
let source_dir = self.repo_root.join(&source.dir);
|
||||
|
||||
match self.install_from_artifact(manifest, force).await {
|
||||
Ok(outcome) => Ok(outcome),
|
||||
@@ -391,6 +441,13 @@ impl RegistryInstaller {
|
||||
let target_dir = match manifest.kind {
|
||||
ManifestKind::Tool => &self.tools_dir,
|
||||
ManifestKind::Channel => &self.channels_dir,
|
||||
ManifestKind::McpServer => {
|
||||
return Err(RegistryError::InvalidManifest {
|
||||
name: manifest.name.clone(),
|
||||
field: "kind",
|
||||
reason: "MCP servers cannot be installed as artifacts".to_string(),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
fs::create_dir_all(target_dir)
|
||||
@@ -458,12 +515,9 @@ impl RegistryInstaller {
|
||||
false
|
||||
}
|
||||
}
|
||||
} else {
|
||||
} else if let Some(ref source) = manifest.source {
|
||||
// Legacy fallback: try source tree
|
||||
let caps_source = self
|
||||
.repo_root
|
||||
.join(&manifest.source.dir)
|
||||
.join(&manifest.source.capabilities);
|
||||
let caps_source = self.repo_root.join(&source.dir).join(&source.capabilities);
|
||||
if caps_source.exists() {
|
||||
fs::copy(&caps_source, &target_caps)
|
||||
.await
|
||||
@@ -472,6 +526,8 @@ impl RegistryInstaller {
|
||||
} else {
|
||||
false
|
||||
}
|
||||
} else {
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
@@ -775,17 +831,19 @@ mod tests {
|
||||
name: name.to_string(),
|
||||
display_name: name.to_string(),
|
||||
kind,
|
||||
version: "0.1.0".to_string(),
|
||||
version: Some("0.1.0".to_string()),
|
||||
description: "test manifest".to_string(),
|
||||
keywords: Vec::new(),
|
||||
source: SourceSpec {
|
||||
source: Some(SourceSpec {
|
||||
dir: source_dir.to_string(),
|
||||
capabilities: format!("{}.capabilities.json", name),
|
||||
crate_name: name.to_string(),
|
||||
},
|
||||
}),
|
||||
artifacts,
|
||||
auth_summary: None,
|
||||
tags: Vec::new(),
|
||||
url: None,
|
||||
auth: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+192
-21
@@ -7,7 +7,7 @@ use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::extensions::{AuthHint, ExtensionKind, ExtensionSource, RegistryEntry};
|
||||
|
||||
/// A single extension manifest loaded from `registry/{tools,channels}/<name>.json`.
|
||||
/// A single extension manifest loaded from `registry/{tools,channels,mcp-servers}/<name>.json`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ExtensionManifest {
|
||||
/// Unique identifier (matches crate name stem, e.g. "slack").
|
||||
@@ -16,11 +16,12 @@ pub struct ExtensionManifest {
|
||||
/// Human-readable name (e.g. "Slack").
|
||||
pub display_name: String,
|
||||
|
||||
/// Whether this is a tool or channel.
|
||||
/// Whether this is a tool, channel, or MCP server.
|
||||
pub kind: ManifestKind,
|
||||
|
||||
/// Semver version from Cargo.toml.
|
||||
pub version: String,
|
||||
/// Semver version from Cargo.toml. Optional for MCP server manifests.
|
||||
#[serde(default)]
|
||||
pub version: Option<String>,
|
||||
|
||||
/// One-line description.
|
||||
pub description: String,
|
||||
@@ -29,8 +30,9 @@ pub struct ExtensionManifest {
|
||||
#[serde(default)]
|
||||
pub keywords: Vec<String>,
|
||||
|
||||
/// Source code location and build info.
|
||||
pub source: SourceSpec,
|
||||
/// Source code location and build info. Absent for MCP server manifests.
|
||||
#[serde(default)]
|
||||
pub source: Option<SourceSpec>,
|
||||
|
||||
/// Pre-built binary artifacts keyed by target triple.
|
||||
#[serde(default)]
|
||||
@@ -43,6 +45,15 @@ pub struct ExtensionManifest {
|
||||
/// Tags for filtering (e.g. "default", "messaging", "google").
|
||||
#[serde(default)]
|
||||
pub tags: Vec<String>,
|
||||
|
||||
/// MCP server URL. Only present for `McpServer` manifests.
|
||||
#[serde(default)]
|
||||
pub url: Option<String>,
|
||||
|
||||
/// MCP auth method: "dcr", "oauth_pre_configured:<setup_url>", or "none".
|
||||
/// Only present for `McpServer` manifests.
|
||||
#[serde(default)]
|
||||
pub auth: Option<String>,
|
||||
}
|
||||
|
||||
/// Extension kind as declared in manifests.
|
||||
@@ -51,6 +62,7 @@ pub struct ExtensionManifest {
|
||||
pub enum ManifestKind {
|
||||
Tool,
|
||||
Channel,
|
||||
McpServer,
|
||||
}
|
||||
|
||||
impl From<ManifestKind> for ExtensionKind {
|
||||
@@ -58,6 +70,7 @@ impl From<ManifestKind> for ExtensionKind {
|
||||
match kind {
|
||||
ManifestKind::Tool => ExtensionKind::WasmTool,
|
||||
ManifestKind::Channel => ExtensionKind::WasmChannel,
|
||||
ManifestKind::McpServer => ExtensionKind::McpServer,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -67,6 +80,7 @@ impl std::fmt::Display for ManifestKind {
|
||||
match self {
|
||||
ManifestKind::Tool => write!(f, "tool"),
|
||||
ManifestKind::Channel => write!(f, "channel"),
|
||||
ManifestKind::McpServer => write!(f, "mcp_server"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -153,12 +167,64 @@ pub struct BundlesFile {
|
||||
impl ExtensionManifest {
|
||||
/// Convert this manifest into a [`RegistryEntry`] for use with the in-chat
|
||||
/// extension discovery system.
|
||||
pub fn to_registry_entry(&self) -> RegistryEntry {
|
||||
let buildable = ExtensionSource::WasmBuildable {
|
||||
source_dir: self.source.dir.clone(),
|
||||
build_dir: Some(self.source.dir.clone()),
|
||||
crate_name: Some(self.source.crate_name.clone()),
|
||||
///
|
||||
/// Returns `None` for MCP server manifests missing a `url` field.
|
||||
pub fn to_registry_entry(&self) -> Option<RegistryEntry> {
|
||||
if self.kind == ManifestKind::McpServer {
|
||||
return self.to_mcp_registry_entry();
|
||||
}
|
||||
|
||||
Some(self.to_wasm_registry_entry())
|
||||
}
|
||||
|
||||
/// Build a [`RegistryEntry`] for an MCP server manifest.
|
||||
fn to_mcp_registry_entry(&self) -> Option<RegistryEntry> {
|
||||
let url = match &self.url {
|
||||
Some(u) => u.clone(),
|
||||
None => {
|
||||
tracing::warn!(
|
||||
"MCP server manifest '{}' is missing 'url' field, skipping",
|
||||
self.name
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let auth_hint = match self.auth.as_deref() {
|
||||
Some("dcr") | None => AuthHint::Dcr,
|
||||
Some("none") => AuthHint::None,
|
||||
Some(other) if other.starts_with("oauth_pre_configured:") => {
|
||||
AuthHint::OAuthPreConfigured {
|
||||
setup_url: other
|
||||
.strip_prefix("oauth_pre_configured:")
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
}
|
||||
}
|
||||
_ => AuthHint::Dcr,
|
||||
};
|
||||
|
||||
Some(RegistryEntry {
|
||||
name: self.name.clone(),
|
||||
display_name: self.display_name.clone(),
|
||||
kind: ExtensionKind::McpServer,
|
||||
description: self.description.clone(),
|
||||
keywords: self.keywords.clone(),
|
||||
source: ExtensionSource::McpUrl { url },
|
||||
fallback_source: None,
|
||||
auth_hint,
|
||||
version: self.version.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a [`RegistryEntry`] for a WASM tool or channel manifest.
|
||||
fn to_wasm_registry_entry(&self) -> RegistryEntry {
|
||||
let source_spec = self.source.as_ref();
|
||||
|
||||
let buildable = source_spec.map(|s| ExtensionSource::WasmBuildable {
|
||||
source_dir: s.dir.clone(),
|
||||
build_dir: Some(s.dir.clone()),
|
||||
crate_name: Some(s.crate_name.clone()),
|
||||
});
|
||||
|
||||
// Prefer pre-built artifact download when a URL is available,
|
||||
// with build-from-source as fallback in case the download fails (e.g., 404).
|
||||
@@ -170,13 +236,32 @@ impl ExtensionManifest {
|
||||
wasm_url: url.clone(),
|
||||
capabilities_url: artifact.capabilities_url.clone(),
|
||||
},
|
||||
Some(Box::new(buildable)),
|
||||
buildable.map(Box::new),
|
||||
)
|
||||
} else if let Some(b) = buildable {
|
||||
(b, None)
|
||||
} else {
|
||||
(buildable, None)
|
||||
// No source spec and no download URL — use a placeholder
|
||||
(
|
||||
ExtensionSource::WasmBuildable {
|
||||
source_dir: String::new(),
|
||||
build_dir: None,
|
||||
crate_name: None,
|
||||
},
|
||||
None,
|
||||
)
|
||||
}
|
||||
} else if let Some(b) = buildable {
|
||||
(b, None)
|
||||
} else {
|
||||
(buildable, None)
|
||||
(
|
||||
ExtensionSource::WasmBuildable {
|
||||
source_dir: String::new(),
|
||||
build_dir: None,
|
||||
crate_name: None,
|
||||
},
|
||||
None,
|
||||
)
|
||||
};
|
||||
|
||||
let auth_hint = match self.auth_summary.as_ref().and_then(|a| a.method.as_deref()) {
|
||||
@@ -195,7 +280,7 @@ impl ExtensionManifest {
|
||||
source,
|
||||
fallback_source,
|
||||
auth_hint,
|
||||
version: Some(self.version.clone()),
|
||||
version: self.version.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -234,10 +319,10 @@ mod tests {
|
||||
let manifest: ExtensionManifest = serde_json::from_str(json).expect("parse manifest");
|
||||
assert_eq!(manifest.name, "slack");
|
||||
assert_eq!(manifest.kind, ManifestKind::Tool);
|
||||
assert_eq!(manifest.version, "0.1.0");
|
||||
assert_eq!(manifest.version.as_deref(), Some("0.1.0"));
|
||||
assert!(manifest.tags.contains(&"default".to_string()));
|
||||
|
||||
let entry = manifest.to_registry_entry();
|
||||
let entry = manifest.to_registry_entry().unwrap();
|
||||
assert_eq!(entry.kind, ExtensionKind::WasmTool);
|
||||
}
|
||||
|
||||
@@ -262,7 +347,7 @@ mod tests {
|
||||
assert!(manifest.auth_summary.is_none());
|
||||
assert!(manifest.artifacts.is_empty());
|
||||
|
||||
let entry = manifest.to_registry_entry();
|
||||
let entry = manifest.to_registry_entry().unwrap();
|
||||
assert_eq!(entry.kind, ExtensionKind::WasmChannel);
|
||||
}
|
||||
|
||||
@@ -296,6 +381,7 @@ mod tests {
|
||||
fn test_manifest_kind_display() {
|
||||
assert_eq!(ManifestKind::Tool.to_string(), "tool");
|
||||
assert_eq!(ManifestKind::Channel.to_string(), "channel");
|
||||
assert_eq!(ManifestKind::McpServer.to_string(), "mcp_server");
|
||||
}
|
||||
|
||||
/// When a manifest has a download URL in artifacts, to_registry_entry()
|
||||
@@ -324,7 +410,7 @@ mod tests {
|
||||
}"#;
|
||||
|
||||
let manifest: ExtensionManifest = serde_json::from_str(json).expect("parse manifest");
|
||||
let entry = manifest.to_registry_entry();
|
||||
let entry = manifest.to_registry_entry().unwrap();
|
||||
|
||||
// Primary source should be WasmDownload
|
||||
assert!(
|
||||
@@ -374,7 +460,7 @@ mod tests {
|
||||
}"#;
|
||||
|
||||
let manifest: ExtensionManifest = serde_json::from_str(json).expect("parse manifest");
|
||||
let entry = manifest.to_registry_entry();
|
||||
let entry = manifest.to_registry_entry().unwrap();
|
||||
|
||||
assert!(
|
||||
matches!(&entry.source, ExtensionSource::WasmBuildable { .. }),
|
||||
@@ -405,7 +491,7 @@ mod tests {
|
||||
}"#;
|
||||
|
||||
let manifest: ExtensionManifest = serde_json::from_str(json).expect("parse manifest");
|
||||
let entry = manifest.to_registry_entry();
|
||||
let entry = manifest.to_registry_entry().unwrap();
|
||||
|
||||
assert!(
|
||||
matches!(&entry.source, ExtensionSource::WasmBuildable { .. }),
|
||||
@@ -416,4 +502,89 @@ mod tests {
|
||||
"Should have no fallback when already using WasmBuildable"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_mcp_server_manifest() {
|
||||
let json = r#"{
|
||||
"name": "notion",
|
||||
"display_name": "Notion",
|
||||
"kind": "mcp_server",
|
||||
"description": "Connect to Notion for reading and writing pages, databases, and comments",
|
||||
"keywords": ["notes", "wiki", "docs", "pages", "database"],
|
||||
"url": "https://mcp.notion.com/mcp",
|
||||
"auth": "dcr"
|
||||
}"#;
|
||||
|
||||
let manifest: ExtensionManifest = serde_json::from_str(json).expect("parse manifest");
|
||||
assert_eq!(manifest.name, "notion");
|
||||
assert_eq!(manifest.kind, ManifestKind::McpServer);
|
||||
assert!(manifest.version.is_none());
|
||||
assert!(manifest.source.is_none());
|
||||
assert_eq!(manifest.url.as_deref(), Some("https://mcp.notion.com/mcp"));
|
||||
assert_eq!(manifest.auth.as_deref(), Some("dcr"));
|
||||
|
||||
let entry = manifest.to_registry_entry().unwrap();
|
||||
assert_eq!(entry.kind, ExtensionKind::McpServer);
|
||||
assert!(
|
||||
matches!(&entry.source, ExtensionSource::McpUrl { url } if url == "https://mcp.notion.com/mcp")
|
||||
);
|
||||
assert!(matches!(&entry.auth_hint, AuthHint::Dcr));
|
||||
assert!(entry.fallback_source.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_server_oauth_pre_configured() {
|
||||
let json = r#"{
|
||||
"name": "custom-mcp",
|
||||
"display_name": "Custom MCP",
|
||||
"kind": "mcp_server",
|
||||
"description": "Custom MCP server",
|
||||
"keywords": [],
|
||||
"url": "https://mcp.example.com",
|
||||
"auth": "oauth_pre_configured:https://example.com/setup"
|
||||
}"#;
|
||||
|
||||
let manifest: ExtensionManifest = serde_json::from_str(json).expect("parse manifest");
|
||||
let entry = manifest.to_registry_entry().unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
&entry.auth_hint,
|
||||
AuthHint::OAuthPreConfigured { setup_url } if setup_url == "https://example.com/setup"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_server_auth_none() {
|
||||
let json = r#"{
|
||||
"name": "local-mcp",
|
||||
"display_name": "Local MCP",
|
||||
"kind": "mcp_server",
|
||||
"description": "Local MCP server",
|
||||
"keywords": [],
|
||||
"url": "http://localhost:8080/mcp",
|
||||
"auth": "none"
|
||||
}"#;
|
||||
|
||||
let manifest: ExtensionManifest = serde_json::from_str(json).expect("parse manifest");
|
||||
let entry = manifest.to_registry_entry().unwrap();
|
||||
|
||||
assert!(matches!(&entry.auth_hint, AuthHint::None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_server_missing_url_returns_none() {
|
||||
let json = r#"{
|
||||
"name": "broken-mcp",
|
||||
"display_name": "Broken MCP",
|
||||
"kind": "mcp_server",
|
||||
"description": "MCP server with no URL",
|
||||
"keywords": []
|
||||
}"#;
|
||||
|
||||
let manifest: ExtensionManifest = serde_json::from_str(json).expect("parse manifest");
|
||||
assert!(
|
||||
manifest.to_registry_entry().is_none(),
|
||||
"MCP manifest without url should return None"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -109,3 +109,59 @@ pub fn create_secrets_store(
|
||||
|
||||
store
|
||||
}
|
||||
|
||||
/// Try to resolve an existing master key from env var or OS keychain.
|
||||
///
|
||||
/// Resolution order:
|
||||
/// 1. `SECRETS_MASTER_KEY` environment variable (hex-encoded)
|
||||
/// 2. OS keychain (macOS Keychain / Linux secret-service)
|
||||
///
|
||||
/// Returns `None` if no key is available (caller should generate one).
|
||||
pub async fn resolve_master_key() -> Option<String> {
|
||||
// 1. Check env var
|
||||
if let Ok(env_key) = std::env::var("SECRETS_MASTER_KEY")
|
||||
&& !env_key.is_empty()
|
||||
{
|
||||
return Some(env_key);
|
||||
}
|
||||
|
||||
// 2. Try OS keychain
|
||||
if let Ok(keychain_key_bytes) = keychain::get_master_key().await {
|
||||
let key_hex: String = keychain_key_bytes
|
||||
.iter()
|
||||
.map(|b| format!("{:02x}", b))
|
||||
.collect();
|
||||
return Some(key_hex);
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Create a `SecretsCrypto` from a master key string.
|
||||
///
|
||||
/// The key is typically hex-encoded (from `generate_master_key_hex` or
|
||||
/// the `SECRETS_MASTER_KEY` env var), but `SecretsCrypto::new` validates
|
||||
/// only key length, not encoding. Any sufficiently long string works.
|
||||
pub fn crypto_from_hex(hex: &str) -> Result<std::sync::Arc<SecretsCrypto>, SecretError> {
|
||||
let crypto = SecretsCrypto::new(secrecy::SecretString::from(hex.to_string()))?;
|
||||
Ok(std::sync::Arc::new(crypto))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_crypto_from_hex_valid() {
|
||||
// 32 bytes = 64 hex chars
|
||||
let hex = "0123456789abcdef".repeat(4); // 64 hex chars
|
||||
let result = crypto_from_hex(&hex);
|
||||
assert!(result.is_ok()); // safety: test assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_crypto_from_hex_invalid() {
|
||||
let result = crypto_from_hex("too_short");
|
||||
assert!(result.is_err()); // safety: test assertion
|
||||
}
|
||||
}
|
||||
|
||||
+555
-7
@@ -220,7 +220,7 @@ pub struct TunnelSettings {
|
||||
}
|
||||
|
||||
/// Channel-specific settings.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChannelSettings {
|
||||
/// Whether HTTP webhook channel is enabled.
|
||||
#[serde(default)]
|
||||
@@ -234,6 +234,30 @@ pub struct ChannelSettings {
|
||||
#[serde(default)]
|
||||
pub http_host: Option<String>,
|
||||
|
||||
/// Whether the web gateway is enabled.
|
||||
#[serde(default = "default_true")]
|
||||
pub gateway_enabled: bool,
|
||||
|
||||
/// Web gateway listen host.
|
||||
#[serde(default)]
|
||||
pub gateway_host: Option<String>,
|
||||
|
||||
/// Web gateway listen port.
|
||||
#[serde(default)]
|
||||
pub gateway_port: Option<u16>,
|
||||
|
||||
/// Web gateway bearer auth token. Auto-generated at gateway startup if unset.
|
||||
#[serde(default)]
|
||||
pub gateway_auth_token: Option<String>,
|
||||
|
||||
/// Web gateway user ID.
|
||||
#[serde(default)]
|
||||
pub gateway_user_id: Option<String>,
|
||||
|
||||
/// Whether the CLI channel is enabled.
|
||||
#[serde(default = "default_true")]
|
||||
pub cli_enabled: bool,
|
||||
|
||||
/// Whether Signal channel is enabled.
|
||||
#[serde(default)]
|
||||
pub signal_enabled: bool,
|
||||
@@ -289,6 +313,34 @@ pub struct ChannelSettings {
|
||||
pub wasm_channels_dir: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl Default for ChannelSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
http_enabled: false,
|
||||
http_port: None,
|
||||
http_host: None,
|
||||
gateway_enabled: true,
|
||||
gateway_host: None,
|
||||
gateway_port: None,
|
||||
gateway_auth_token: None,
|
||||
gateway_user_id: None,
|
||||
cli_enabled: true,
|
||||
signal_enabled: false,
|
||||
signal_http_url: None,
|
||||
signal_account: None,
|
||||
signal_allow_from: None,
|
||||
signal_allow_from_groups: None,
|
||||
signal_dm_policy: None,
|
||||
signal_group_policy: None,
|
||||
signal_group_allow_from: None,
|
||||
wasm_channel_owner_ids: std::collections::HashMap::new(),
|
||||
wasm_channels: Vec::new(),
|
||||
wasm_channels_enabled: true,
|
||||
wasm_channels_dir: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Heartbeat configuration.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct HeartbeatSettings {
|
||||
@@ -837,19 +889,16 @@ impl Settings {
|
||||
.map_err(|e| format!("Failed to serialize settings: {}", e))?;
|
||||
|
||||
let parts: Vec<&str> = path.split('.').collect();
|
||||
if parts.is_empty() {
|
||||
return Err("Empty path".to_string());
|
||||
}
|
||||
let (final_key, parent_parts) =
|
||||
parts.split_last().ok_or_else(|| "Empty path".to_string())?;
|
||||
|
||||
// Navigate to parent and set the final key
|
||||
let mut current = &mut json;
|
||||
for part in &parts[..parts.len() - 1] {
|
||||
for part in parent_parts {
|
||||
current = current
|
||||
.get_mut(*part)
|
||||
.ok_or_else(|| format!("Path not found: {}", path))?;
|
||||
}
|
||||
|
||||
let final_key = parts.last().unwrap();
|
||||
let obj = current
|
||||
.as_object_mut()
|
||||
.ok_or_else(|| format!("Parent is not an object: {}", path))?;
|
||||
@@ -1698,4 +1747,503 @@ mod tests {
|
||||
"None selected_model should stay None"
|
||||
);
|
||||
}
|
||||
|
||||
// === Wizard re-run regression tests ===
|
||||
//
|
||||
// These tests simulate the merge ordering used by the wizard's `run()` method
|
||||
// to verify that re-running the wizard (or a subset of steps) doesn't
|
||||
// accidentally reset settings from prior runs.
|
||||
|
||||
/// Simulates `ironclaw onboard --provider-only` re-running on a fully
|
||||
/// configured installation. Only provider + model should change; all
|
||||
/// other settings (channels, embeddings, heartbeat) must survive.
|
||||
#[test]
|
||||
fn provider_only_rerun_preserves_unrelated_settings() {
|
||||
// Prior completed run with everything configured
|
||||
let prior = Settings {
|
||||
onboard_completed: true,
|
||||
database_backend: Some("libsql".to_string()),
|
||||
libsql_path: Some("/home/user/.ironclaw/ironclaw.db".to_string()),
|
||||
llm_backend: Some("openai".to_string()),
|
||||
selected_model: Some("gpt-4o".to_string()),
|
||||
embeddings: EmbeddingsSettings {
|
||||
enabled: true,
|
||||
provider: "openai".to_string(),
|
||||
model: "text-embedding-3-small".to_string(),
|
||||
},
|
||||
channels: ChannelSettings {
|
||||
http_enabled: true,
|
||||
http_port: Some(8080),
|
||||
signal_enabled: true,
|
||||
signal_account: Some("+1234567890".to_string()),
|
||||
wasm_channels: vec!["telegram".to_string()],
|
||||
..Default::default()
|
||||
},
|
||||
heartbeat: HeartbeatSettings {
|
||||
enabled: true,
|
||||
interval_secs: 900,
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let db_map = prior.to_db_map();
|
||||
|
||||
// provider_only mode: reconnect_existing_db loads from DB,
|
||||
// then user picks a new provider + model via step_inference_provider
|
||||
let mut current = Settings::from_db_map(&db_map);
|
||||
|
||||
// Simulate step_inference_provider: user switches to anthropic
|
||||
current.llm_backend = Some("anthropic".to_string());
|
||||
current.selected_model = None; // cleared because backend changed
|
||||
|
||||
// Simulate step_model_selection: user picks a model
|
||||
current.selected_model = Some("claude-sonnet-4-5".to_string());
|
||||
|
||||
// Verify: provider/model changed
|
||||
assert_eq!(current.llm_backend.as_deref(), Some("anthropic"));
|
||||
assert_eq!(current.selected_model.as_deref(), Some("claude-sonnet-4-5"));
|
||||
|
||||
// Verify: everything else preserved
|
||||
assert!(current.channels.http_enabled, "HTTP channel must survive");
|
||||
assert_eq!(current.channels.http_port, Some(8080));
|
||||
assert!(current.channels.signal_enabled, "Signal must survive");
|
||||
assert_eq!(
|
||||
current.channels.wasm_channels,
|
||||
vec!["telegram".to_string()],
|
||||
"WASM channels must survive"
|
||||
);
|
||||
assert!(current.embeddings.enabled, "Embeddings must survive");
|
||||
assert_eq!(current.embeddings.provider, "openai");
|
||||
assert!(current.heartbeat.enabled, "Heartbeat must survive");
|
||||
assert_eq!(current.heartbeat.interval_secs, 900);
|
||||
assert_eq!(
|
||||
current.database_backend.as_deref(),
|
||||
Some("libsql"),
|
||||
"DB backend must survive"
|
||||
);
|
||||
}
|
||||
|
||||
/// Simulates `ironclaw onboard --channels-only` re-running on a fully
|
||||
/// configured installation. Only channel settings should change;
|
||||
/// provider, model, embeddings, heartbeat must survive.
|
||||
#[test]
|
||||
fn channels_only_rerun_preserves_unrelated_settings() {
|
||||
let prior = Settings {
|
||||
onboard_completed: true,
|
||||
database_backend: Some("postgres".to_string()),
|
||||
database_url: Some("postgres://host/db".to_string()),
|
||||
llm_backend: Some("anthropic".to_string()),
|
||||
selected_model: Some("claude-sonnet-4-5".to_string()),
|
||||
embeddings: EmbeddingsSettings {
|
||||
enabled: true,
|
||||
provider: "nearai".to_string(),
|
||||
model: "text-embedding-3-small".to_string(),
|
||||
},
|
||||
heartbeat: HeartbeatSettings {
|
||||
enabled: true,
|
||||
interval_secs: 1800,
|
||||
..Default::default()
|
||||
},
|
||||
channels: ChannelSettings {
|
||||
http_enabled: false,
|
||||
wasm_channels: vec!["telegram".to_string()],
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let db_map = prior.to_db_map();
|
||||
|
||||
// channels_only mode: reconnect_existing_db loads from DB
|
||||
let mut current = Settings::from_db_map(&db_map);
|
||||
|
||||
// Simulate step_channels: user enables HTTP and adds discord
|
||||
current.channels.http_enabled = true;
|
||||
current.channels.http_port = Some(9090);
|
||||
current.channels.wasm_channels = vec!["telegram".to_string(), "discord".to_string()];
|
||||
|
||||
// Verify: channels changed
|
||||
assert!(current.channels.http_enabled);
|
||||
assert_eq!(current.channels.http_port, Some(9090));
|
||||
assert_eq!(current.channels.wasm_channels.len(), 2);
|
||||
|
||||
// Verify: everything else preserved
|
||||
assert_eq!(current.llm_backend.as_deref(), Some("anthropic"));
|
||||
assert_eq!(current.selected_model.as_deref(), Some("claude-sonnet-4-5"));
|
||||
assert!(current.embeddings.enabled);
|
||||
assert_eq!(current.embeddings.provider, "nearai");
|
||||
assert!(current.heartbeat.enabled);
|
||||
assert_eq!(current.heartbeat.interval_secs, 1800);
|
||||
}
|
||||
|
||||
/// Simulates quick mode re-run on an installation that previously
|
||||
/// completed a full setup. Quick mode only touches DB + security +
|
||||
/// provider + model; channels, embeddings, heartbeat, extensions
|
||||
/// should survive via the merge_from ordering.
|
||||
#[test]
|
||||
fn quick_mode_rerun_preserves_prior_channels_and_heartbeat() {
|
||||
let prior = Settings {
|
||||
onboard_completed: true,
|
||||
database_backend: Some("libsql".to_string()),
|
||||
libsql_path: Some("/home/user/.ironclaw/ironclaw.db".to_string()),
|
||||
llm_backend: Some("openai".to_string()),
|
||||
selected_model: Some("gpt-4o".to_string()),
|
||||
channels: ChannelSettings {
|
||||
http_enabled: true,
|
||||
http_port: Some(8080),
|
||||
signal_enabled: true,
|
||||
wasm_channels: vec!["telegram".to_string()],
|
||||
..Default::default()
|
||||
},
|
||||
embeddings: EmbeddingsSettings {
|
||||
enabled: true,
|
||||
provider: "openai".to_string(),
|
||||
model: "text-embedding-3-small".to_string(),
|
||||
},
|
||||
heartbeat: HeartbeatSettings {
|
||||
enabled: true,
|
||||
interval_secs: 600,
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let db_map = prior.to_db_map();
|
||||
let from_db = Settings::from_db_map(&db_map);
|
||||
|
||||
// Quick mode flow:
|
||||
// 1. auto_setup_database sets DB fields
|
||||
let step1 = Settings {
|
||||
database_backend: Some("libsql".to_string()),
|
||||
libsql_path: Some("/home/user/.ironclaw/ironclaw.db".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// 2. try_load_existing_settings → merge DB → merge step1 on top
|
||||
let mut current = step1.clone();
|
||||
current.merge_from(&from_db);
|
||||
current.merge_from(&step1);
|
||||
|
||||
// 3. step_inference_provider: user picks anthropic this time
|
||||
current.llm_backend = Some("anthropic".to_string());
|
||||
current.selected_model = None; // cleared because backend changed
|
||||
|
||||
// 4. step_model_selection: user picks model
|
||||
current.selected_model = Some("claude-opus-4-6".to_string());
|
||||
|
||||
// Verify: provider/model updated
|
||||
assert_eq!(current.llm_backend.as_deref(), Some("anthropic"));
|
||||
assert_eq!(current.selected_model.as_deref(), Some("claude-opus-4-6"));
|
||||
|
||||
// Verify: channels, embeddings, heartbeat survived quick mode
|
||||
assert!(
|
||||
current.channels.http_enabled,
|
||||
"HTTP channel must survive quick mode re-run"
|
||||
);
|
||||
assert_eq!(current.channels.http_port, Some(8080));
|
||||
assert!(
|
||||
current.channels.signal_enabled,
|
||||
"Signal must survive quick mode re-run"
|
||||
);
|
||||
assert_eq!(
|
||||
current.channels.wasm_channels,
|
||||
vec!["telegram".to_string()],
|
||||
"WASM channels must survive quick mode re-run"
|
||||
);
|
||||
assert!(
|
||||
current.embeddings.enabled,
|
||||
"Embeddings must survive quick mode re-run"
|
||||
);
|
||||
assert!(
|
||||
current.heartbeat.enabled,
|
||||
"Heartbeat must survive quick mode re-run"
|
||||
);
|
||||
assert_eq!(current.heartbeat.interval_secs, 600);
|
||||
}
|
||||
|
||||
/// Full wizard re-run where user keeps the same provider. The model
|
||||
/// selection from the prior run should be pre-populated (not reset).
|
||||
///
|
||||
/// Regression: re-running with the same provider should preserve model.
|
||||
#[test]
|
||||
fn full_rerun_same_provider_preserves_model_through_merge() {
|
||||
let prior = Settings {
|
||||
onboard_completed: true,
|
||||
database_backend: Some("postgres".to_string()),
|
||||
database_url: Some("postgres://host/db".to_string()),
|
||||
llm_backend: Some("anthropic".to_string()),
|
||||
selected_model: Some("claude-sonnet-4-5".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let db_map = prior.to_db_map();
|
||||
let from_db = Settings::from_db_map(&db_map);
|
||||
|
||||
// Step 1: user keeps same DB
|
||||
let step1 = Settings {
|
||||
database_backend: Some("postgres".to_string()),
|
||||
database_url: Some("postgres://host/db".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut current = step1.clone();
|
||||
current.merge_from(&from_db);
|
||||
current.merge_from(&step1);
|
||||
|
||||
// After merge, prior settings recovered
|
||||
assert_eq!(
|
||||
current.llm_backend.as_deref(),
|
||||
Some("anthropic"),
|
||||
"Prior provider must be recovered from DB"
|
||||
);
|
||||
assert_eq!(
|
||||
current.selected_model.as_deref(),
|
||||
Some("claude-sonnet-4-5"),
|
||||
"Prior model must be recovered from DB"
|
||||
);
|
||||
|
||||
// Step 3: user picks same provider (anthropic)
|
||||
// set_llm_backend_preserving_model checks if backend changed
|
||||
let backend_changed = current.llm_backend.as_deref() != Some("anthropic");
|
||||
current.llm_backend = Some("anthropic".to_string());
|
||||
if backend_changed {
|
||||
current.selected_model = None;
|
||||
}
|
||||
|
||||
// Model should NOT be cleared since backend didn't change
|
||||
assert_eq!(
|
||||
current.selected_model.as_deref(),
|
||||
Some("claude-sonnet-4-5"),
|
||||
"Model must survive when re-selecting same provider"
|
||||
);
|
||||
}
|
||||
|
||||
/// Full wizard re-run where user switches provider. Model should be
|
||||
/// cleared since the old model is invalid for the new backend.
|
||||
#[test]
|
||||
fn full_rerun_different_provider_clears_model_through_merge() {
|
||||
let prior = Settings {
|
||||
onboard_completed: true,
|
||||
database_backend: Some("postgres".to_string()),
|
||||
database_url: Some("postgres://host/db".to_string()),
|
||||
llm_backend: Some("anthropic".to_string()),
|
||||
selected_model: Some("claude-sonnet-4-5".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let db_map = prior.to_db_map();
|
||||
let from_db = Settings::from_db_map(&db_map);
|
||||
|
||||
// Step 1 merge
|
||||
let step1 = Settings {
|
||||
database_backend: Some("postgres".to_string()),
|
||||
database_url: Some("postgres://host/db".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let mut current = step1.clone();
|
||||
current.merge_from(&from_db);
|
||||
current.merge_from(&step1);
|
||||
|
||||
// Step 3: user switches to openai
|
||||
let backend_changed = current.llm_backend.as_deref() != Some("openai");
|
||||
assert!(backend_changed, "switching providers should be detected");
|
||||
current.llm_backend = Some("openai".to_string());
|
||||
if backend_changed {
|
||||
current.selected_model = None;
|
||||
}
|
||||
|
||||
assert_eq!(current.llm_backend.as_deref(), Some("openai"));
|
||||
assert!(
|
||||
current.selected_model.is_none(),
|
||||
"Model must be cleared when switching providers"
|
||||
);
|
||||
}
|
||||
|
||||
/// Simulates incremental save correctness: persist_after_step after
|
||||
/// Step 3 (provider) should not clobber settings set in Step 2 (security).
|
||||
///
|
||||
/// The wizard persists the full settings object after each step. This
|
||||
/// test verifies that incremental saves are idempotent for prior steps.
|
||||
#[test]
|
||||
fn incremental_persist_does_not_clobber_prior_steps() {
|
||||
// After steps 1-2, settings has DB + security
|
||||
let after_step2 = Settings {
|
||||
database_backend: Some("libsql".to_string()),
|
||||
secrets_master_key_source: KeySource::Keychain,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// persist_after_step saves to DB
|
||||
let db_map_after_step2 = after_step2.to_db_map();
|
||||
|
||||
// Step 3 adds provider
|
||||
let mut after_step3 = after_step2.clone();
|
||||
after_step3.llm_backend = Some("openai".to_string());
|
||||
|
||||
// persist_after_step saves again — the full settings object
|
||||
let db_map_after_step3 = after_step3.to_db_map();
|
||||
|
||||
// Reload from DB after step 3
|
||||
let restored = Settings::from_db_map(&db_map_after_step3);
|
||||
|
||||
// Step 2's settings must survive step 3's persist
|
||||
assert_eq!(
|
||||
restored.secrets_master_key_source,
|
||||
KeySource::Keychain,
|
||||
"Step 2 security setting must survive step 3 persist"
|
||||
);
|
||||
assert_eq!(
|
||||
restored.database_backend.as_deref(),
|
||||
Some("libsql"),
|
||||
"Step 1 DB setting must survive step 3 persist"
|
||||
);
|
||||
assert_eq!(
|
||||
restored.llm_backend.as_deref(),
|
||||
Some("openai"),
|
||||
"Step 3 provider setting must be saved"
|
||||
);
|
||||
|
||||
// Also verify that a partial step 2 reload doesn't regress
|
||||
// (loading the step 2 snapshot and merging with step 3 state)
|
||||
let from_step2_db = Settings::from_db_map(&db_map_after_step2);
|
||||
let mut merged = after_step3.clone();
|
||||
merged.merge_from(&from_step2_db);
|
||||
|
||||
assert_eq!(
|
||||
merged.llm_backend.as_deref(),
|
||||
Some("openai"),
|
||||
"Step 3 provider must not be clobbered by step 2 snapshot merge"
|
||||
);
|
||||
assert_eq!(
|
||||
merged.secrets_master_key_source,
|
||||
KeySource::Keychain,
|
||||
"Step 2 security must survive merge"
|
||||
);
|
||||
}
|
||||
|
||||
/// Switching database backend should allow fresh connection settings.
|
||||
/// When user switches from postgres to libsql, the old database_url
|
||||
/// should not prevent the new libsql_path from being used.
|
||||
#[test]
|
||||
fn switching_db_backend_allows_fresh_connection_settings() {
|
||||
let prior = Settings {
|
||||
database_backend: Some("postgres".to_string()),
|
||||
database_url: Some("postgres://host/db".to_string()),
|
||||
llm_backend: Some("openai".to_string()),
|
||||
selected_model: Some("gpt-4o".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let db_map = prior.to_db_map();
|
||||
let from_db = Settings::from_db_map(&db_map);
|
||||
|
||||
// User picks libsql this time, wizard clears stale postgres settings
|
||||
let step1 = Settings {
|
||||
database_backend: Some("libsql".to_string()),
|
||||
libsql_path: Some("/home/user/.ironclaw/ironclaw.db".to_string()),
|
||||
database_url: None, // explicitly not set for libsql
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut current = step1.clone();
|
||||
current.merge_from(&from_db);
|
||||
current.merge_from(&step1);
|
||||
|
||||
// libsql chosen
|
||||
assert_eq!(current.database_backend.as_deref(), Some("libsql"));
|
||||
assert_eq!(
|
||||
current.libsql_path.as_deref(),
|
||||
Some("/home/user/.ironclaw/ironclaw.db")
|
||||
);
|
||||
|
||||
// Prior provider/model should survive (unrelated to DB switch)
|
||||
assert_eq!(current.llm_backend.as_deref(), Some("openai"));
|
||||
assert_eq!(current.selected_model.as_deref(), Some("gpt-4o"));
|
||||
|
||||
// Note: database_url from prior run persists in merge because
|
||||
// step1.database_url is None (== default), so merge_from doesn't
|
||||
// override it. This is expected — the .env writer decides which
|
||||
// vars to emit based on database_backend. The stale URL is
|
||||
// harmless because the libsql backend ignores it.
|
||||
assert_eq!(
|
||||
current.database_url.as_deref(),
|
||||
Some("postgres://host/db"),
|
||||
"stale database_url persists (harmless, ignored by libsql backend)"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression: merge_from must handle boolean fields correctly.
|
||||
/// A prior run with heartbeat.enabled=true must not be reset to false
|
||||
/// when merging with a Settings that has heartbeat.enabled=false (default).
|
||||
#[test]
|
||||
fn merge_preserves_true_booleans_when_overlay_has_default_false() {
|
||||
let prior = Settings {
|
||||
heartbeat: HeartbeatSettings {
|
||||
enabled: true,
|
||||
interval_secs: 600,
|
||||
..Default::default()
|
||||
},
|
||||
channels: ChannelSettings {
|
||||
http_enabled: true,
|
||||
signal_enabled: true,
|
||||
..Default::default()
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let db_map = prior.to_db_map();
|
||||
let from_db = Settings::from_db_map(&db_map);
|
||||
|
||||
// New wizard run only sets DB (everything else is default/false)
|
||||
let step1 = Settings {
|
||||
database_backend: Some("libsql".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut current = step1.clone();
|
||||
current.merge_from(&from_db);
|
||||
current.merge_from(&step1);
|
||||
|
||||
// true booleans from prior run must survive
|
||||
assert!(
|
||||
current.heartbeat.enabled,
|
||||
"heartbeat.enabled=true must not be reset to false by default overlay"
|
||||
);
|
||||
assert!(
|
||||
current.channels.http_enabled,
|
||||
"http_enabled=true must not be reset to false by default overlay"
|
||||
);
|
||||
assert!(
|
||||
current.channels.signal_enabled,
|
||||
"signal_enabled=true must not be reset to false by default overlay"
|
||||
);
|
||||
assert_eq!(current.heartbeat.interval_secs, 600);
|
||||
}
|
||||
|
||||
/// Regression: embeddings settings (provider, model, enabled) must
|
||||
/// survive a wizard re-run that doesn't touch step 5.
|
||||
#[test]
|
||||
fn embeddings_survive_rerun_that_skips_step5() {
|
||||
let prior = Settings {
|
||||
onboard_completed: true,
|
||||
llm_backend: Some("nearai".to_string()),
|
||||
selected_model: Some("qwen".to_string()),
|
||||
embeddings: EmbeddingsSettings {
|
||||
enabled: true,
|
||||
provider: "nearai".to_string(),
|
||||
model: "text-embedding-3-large".to_string(),
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let db_map = prior.to_db_map();
|
||||
let from_db = Settings::from_db_map(&db_map);
|
||||
|
||||
// Full re-run: step 1 only sets DB
|
||||
let step1 = Settings {
|
||||
database_backend: Some("libsql".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let mut current = step1.clone();
|
||||
current.merge_from(&from_db);
|
||||
current.merge_from(&step1);
|
||||
|
||||
// Before step 5 (embeddings) runs, check that prior values are present
|
||||
assert!(current.embeddings.enabled);
|
||||
assert_eq!(current.embeddings.provider, "nearai");
|
||||
assert_eq!(current.embeddings.model, "text-embedding-3-large");
|
||||
}
|
||||
}
|
||||
|
||||
+17
-13
@@ -114,6 +114,13 @@ Step 9: Background Tasks (heartbeat)
|
||||
|
||||
**Goal:** Select backend, establish connection, run migrations.
|
||||
|
||||
**Init delegation:** Backend-specific connection logic lives in `src/db/mod.rs`
|
||||
(`connect_without_migrations()`), not in the wizard. The wizard calls
|
||||
`test_database_connection()` which delegates to the db module factory. Feature-flag
|
||||
branching (`#[cfg(feature = ...)]`) is confined to `src/db/mod.rs`. PostgreSQL
|
||||
validation (version >= 15, pgvector) is handled by `validate_postgres()` in
|
||||
`src/db/mod.rs`.
|
||||
|
||||
**Decision tree:**
|
||||
|
||||
```
|
||||
@@ -121,26 +128,23 @@ Both features compiled?
|
||||
├─ Yes → DATABASE_BACKEND env var set?
|
||||
│ ├─ Yes → use that backend
|
||||
│ └─ No → interactive selection (PostgreSQL vs libSQL)
|
||||
├─ Only postgres feature → step_database_postgres()
|
||||
└─ Only libsql feature → step_database_libsql()
|
||||
├─ Only postgres feature → prompt for DATABASE_URL, test connection
|
||||
└─ Only libsql feature → prompt for path, test connection
|
||||
```
|
||||
|
||||
**PostgreSQL path** (`step_database_postgres`):
|
||||
**PostgreSQL path:**
|
||||
1. Check `DATABASE_URL` from env or settings
|
||||
2. Test connection (creates `deadpool_postgres::Pool`)
|
||||
3. Optionally run refinery migrations
|
||||
4. Store pool in `self.db_pool`
|
||||
2. Test connection via `connect_without_migrations()` (validates version, pgvector)
|
||||
3. Optionally run migrations
|
||||
|
||||
**libSQL path** (`step_database_libsql`):
|
||||
**libSQL path:**
|
||||
1. Offer local path (default: `~/.ironclaw/ironclaw.db`)
|
||||
2. Optional Turso cloud sync (URL + auth token)
|
||||
3. Test connection (creates `LibSqlBackend`)
|
||||
3. Test connection via `connect_without_migrations()`
|
||||
4. Always run migrations (idempotent CREATE IF NOT EXISTS)
|
||||
5. Store backend in `self.db_backend`
|
||||
|
||||
**Invariant:** After Step 1, exactly one of `self.db_pool` or
|
||||
`self.db_backend` is `Some`. This is required for settings persistence
|
||||
in `save_and_summarize()`.
|
||||
**Invariant:** After Step 1, `self.db` is `Some(Arc<dyn Database>)`.
|
||||
This is required for settings persistence in `save_and_summarize()`.
|
||||
|
||||
---
|
||||
|
||||
@@ -338,7 +342,7 @@ key first, then falls back to the standard env var.
|
||||
1. Check `self.secrets_crypto` (set in Step 2) → use if available
|
||||
2. Else try `SECRETS_MASTER_KEY` env var
|
||||
3. Else try `get_master_key()` from keychain (only in `channels_only` mode)
|
||||
4. Create backend-appropriate secrets store (respects selected database backend)
|
||||
4. Create secrets store using `self.db` (`Arc<dyn Database>`)
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1016,7 +1016,7 @@ fn validation_placeholder_regex() -> &'static regex::Regex {
|
||||
static PLACEHOLDER_RE: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
|
||||
PLACEHOLDER_RE.get_or_init(|| {
|
||||
regex::Regex::new(r"\{([A-Za-z0-9_]+)\}")
|
||||
.expect("validation placeholder regex must compile")
|
||||
.expect("validation placeholder regex must compile") // safety: hardcoded literal
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+253
-978
File diff suppressed because it is too large
Load Diff
+4
-4
@@ -48,7 +48,7 @@ pub const MAX_PROMPT_FILE_SIZE: u64 = 64 * 1024;
|
||||
|
||||
/// Regex for validating skill names: alphanumeric, hyphens, underscores, dots.
|
||||
static SKILL_NAME_PATTERN: std::sync::LazyLock<Regex> =
|
||||
std::sync::LazyLock::new(|| Regex::new(r"^[a-zA-Z0-9][a-zA-Z0-9._-]{0,63}$").unwrap());
|
||||
std::sync::LazyLock::new(|| Regex::new(r"^[a-zA-Z0-9][a-zA-Z0-9._-]{0,63}$").unwrap()); // safety: hardcoded literal
|
||||
|
||||
/// Validate a skill name against the allowed pattern.
|
||||
pub fn validate_skill_name(name: &str) -> bool {
|
||||
@@ -268,13 +268,13 @@ pub fn escape_skill_content(content: &str) -> String {
|
||||
// Match `<` followed by optional `/`, optional whitespace/control chars,
|
||||
// then `skill` (case-insensitive). Catches both opening and closing tags:
|
||||
// `<skill`, `</skill`, `< skill`, `</\0skill`, `<SKILL`, etc.
|
||||
Regex::new(r"(?i)</?[\s\x00]*skill").unwrap()
|
||||
Regex::new(r"(?i)</?[\s\x00]*skill").unwrap() // safety: hardcoded literal
|
||||
});
|
||||
|
||||
SKILL_TAG_RE
|
||||
.replace_all(content, |caps: ®ex::Captures| {
|
||||
// Replace leading `<` with `<` to neutralize the tag
|
||||
let matched = caps.get(0).unwrap().as_str();
|
||||
// Replace leading `<` with `<` to neutralize the tag.
|
||||
let matched = caps.get(0).unwrap().as_str(); // safety: group 0 always exists
|
||||
format!("<{}", &matched[1..])
|
||||
})
|
||||
.into_owned()
|
||||
|
||||
@@ -43,8 +43,8 @@ use crate::error::ToolError as AgentToolError;
|
||||
use crate::llm::{
|
||||
ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolDefinition,
|
||||
};
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
|
||||
use crate::tools::{ToolRegistry, prepare_tool_params};
|
||||
|
||||
/// Requirement specification for building software.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -776,10 +776,11 @@ Create alongside the .wasm file to grant capabilities:
|
||||
self.tools.get(tool_name).await.ok_or_else(|| {
|
||||
ToolError::ExecutionFailed(format!("Tool not found: {}", tool_name))
|
||||
})?;
|
||||
let normalized_params = prepare_tool_params(tool.as_ref(), params);
|
||||
|
||||
// Execute with a dummy context (build tools don't need job context)
|
||||
let ctx = JobContext::default();
|
||||
tool.execute(params.clone(), &ctx).await
|
||||
tool.execute(normalized_params, &ctx).await
|
||||
}
|
||||
|
||||
/// Find the build artifact based on project type.
|
||||
|
||||
@@ -330,7 +330,11 @@ impl CreateJobTool {
|
||||
ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let start = std::time::Instant::now();
|
||||
let jm = self.job_manager.as_ref().expect("sandbox deps required");
|
||||
let jm = self.job_manager.as_ref().ok_or_else(|| {
|
||||
ToolError::ExecutionFailed(
|
||||
"Sandbox execution requires a configured job manager (container runtime not available)".to_string(),
|
||||
)
|
||||
})?;
|
||||
|
||||
let job_id = Uuid::new_v4();
|
||||
let (project_dir, browse_id) = resolve_project_dir(explicit_dir, job_id)?;
|
||||
@@ -411,7 +415,19 @@ impl CreateJobTool {
|
||||
// loop stops consuming from inject_tx the send will fail and the
|
||||
// monitor terminates. No JoinHandle is retained.
|
||||
if let (Some(etx), Some(itx)) = (&self.event_tx, &self.inject_tx) {
|
||||
crate::agent::job_monitor::spawn_job_monitor(job_id, etx.subscribe(), itx.clone());
|
||||
if let Some(route) = monitor_route_from_ctx(ctx) {
|
||||
crate::agent::job_monitor::spawn_job_monitor(
|
||||
job_id,
|
||||
etx.subscribe(),
|
||||
itx.clone(),
|
||||
route,
|
||||
);
|
||||
} else {
|
||||
tracing::debug!(
|
||||
job_id = %job_id,
|
||||
"Skipping job monitor injection due to missing route metadata"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let result = serde_json::json!({
|
||||
@@ -676,6 +692,36 @@ fn resolve_project_dir(
|
||||
Ok((canonical_dir, browse_id))
|
||||
}
|
||||
|
||||
fn monitor_route_from_ctx(ctx: &JobContext) -> Option<crate::agent::job_monitor::JobMonitorRoute> {
|
||||
// notify_channel is required — without it we don't know which channel to
|
||||
// route the monitor output to, so return None to skip monitoring entirely.
|
||||
let channel = ctx
|
||||
.metadata
|
||||
.get("notify_channel")
|
||||
.and_then(|v| v.as_str())?
|
||||
.to_string();
|
||||
// notify_user is optional — fall back to the job's own user_id, which is
|
||||
// always present. The channel is the routing decision; the user is just
|
||||
// for attribution and can default safely.
|
||||
let user_id = ctx
|
||||
.metadata
|
||||
.get("notify_user")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or(&ctx.user_id)
|
||||
.to_string();
|
||||
let thread_id = ctx
|
||||
.metadata
|
||||
.get("notify_thread_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
Some(crate::agent::job_monitor::JobMonitorRoute {
|
||||
channel,
|
||||
user_id,
|
||||
thread_id,
|
||||
})
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for CreateJobTool {
|
||||
fn name(&self) -> &str {
|
||||
@@ -1379,6 +1425,31 @@ mod tests {
|
||||
assert_eq!(tool.execution_timeout(), Duration::from_secs(30));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sandbox_without_job_manager_returns_error() {
|
||||
let manager = Arc::new(ContextManager::new(5));
|
||||
// Create tool without sandbox deps — job_manager is None.
|
||||
let tool = CreateJobTool::new(manager);
|
||||
assert!(!tool.sandbox_enabled());
|
||||
|
||||
let result = tool
|
||||
.execute_sandbox(
|
||||
"test task",
|
||||
None,
|
||||
false,
|
||||
JobMode::Worker,
|
||||
vec![],
|
||||
&JobContext::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let err = result.unwrap_err();
|
||||
assert!(
|
||||
matches!(err, ToolError::ExecutionFailed(_)),
|
||||
"expected ExecutionFailed, got: {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_jobs_tool() {
|
||||
let manager = Arc::new(ContextManager::new(5));
|
||||
|
||||
@@ -301,7 +301,11 @@ impl Tool for SkillInstallTool {
|
||||
let content = if let Some(raw) = params.get("content").and_then(|v| v.as_str()) {
|
||||
// Direct content provided
|
||||
raw.to_string()
|
||||
} else if let Some(url) = params.get("url").and_then(|v| v.as_str()) {
|
||||
} else if let Some(url) = params
|
||||
.get("url")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|s| !s.is_empty())
|
||||
{
|
||||
// Fetch from explicit URL
|
||||
fetch_skill_content(url).await?
|
||||
} else {
|
||||
@@ -1297,4 +1301,23 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_url_param_is_treated_as_absent() {
|
||||
// LLMs sometimes pass "" for optional parameters instead of omitting them.
|
||||
// Before the fix, url: "" would match Some("") and attempt to fetch from an
|
||||
// empty URL (failing with an invalid URL error) instead of falling through to
|
||||
// the catalog lookup. The full execute path cannot be tested here without a
|
||||
// real catalog and database, so this test verifies the parameter filtering
|
||||
// behaviour directly.
|
||||
let params = serde_json::json!({"name": "my-skill", "url": ""});
|
||||
let url = params
|
||||
.get("url")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|s| !s.is_empty());
|
||||
assert!(
|
||||
url.is_none(),
|
||||
"empty url string should be treated as absent"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,367 @@
|
||||
pub(crate) fn prepare_tool_params(
|
||||
tool: &dyn crate::tools::tool::Tool,
|
||||
params: &serde_json::Value,
|
||||
) -> serde_json::Value {
|
||||
prepare_params_for_schema(params, &tool.discovery_schema())
|
||||
}
|
||||
|
||||
pub(crate) fn prepare_params_for_schema(
|
||||
params: &serde_json::Value,
|
||||
schema: &serde_json::Value,
|
||||
) -> serde_json::Value {
|
||||
coerce_value(params, schema)
|
||||
}
|
||||
|
||||
fn coerce_value(value: &serde_json::Value, schema: &serde_json::Value) -> serde_json::Value {
|
||||
// This coercer intentionally handles the concrete schema shapes we expose in
|
||||
// discovery today. It does not resolve combinators like anyOf/oneOf/allOf or
|
||||
// references via $ref; those schemas pass through unchanged unless they also
|
||||
// advertise a directly coercible type/property shape.
|
||||
if value.is_null() {
|
||||
return value.clone();
|
||||
}
|
||||
|
||||
if let Some(s) = value.as_str() {
|
||||
return coerce_string_value(s, schema).unwrap_or_else(|| value.clone());
|
||||
}
|
||||
|
||||
if let Some(items) = value.as_array() {
|
||||
if !schema_allows_type(schema, "array") {
|
||||
return value.clone();
|
||||
}
|
||||
|
||||
let Some(item_schema) = schema.get("items") else {
|
||||
return value.clone();
|
||||
};
|
||||
|
||||
return serde_json::Value::Array(
|
||||
items
|
||||
.iter()
|
||||
.map(|item| coerce_value(item, item_schema))
|
||||
.collect(),
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(obj) = value.as_object() {
|
||||
if !schema_allows_type(schema, "object") {
|
||||
return value.clone();
|
||||
}
|
||||
|
||||
let properties = schema.get("properties").and_then(|p| p.as_object());
|
||||
let additional_schema = schema.get("additionalProperties").filter(|v| v.is_object());
|
||||
let mut coerced = obj.clone();
|
||||
|
||||
for (key, current) in &mut coerced {
|
||||
if let Some(prop_schema) = properties.and_then(|props| props.get(key)) {
|
||||
*current = coerce_value(current, prop_schema);
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(additional_schema) = additional_schema {
|
||||
*current = coerce_value(current, additional_schema);
|
||||
}
|
||||
}
|
||||
|
||||
return serde_json::Value::Object(coerced);
|
||||
}
|
||||
|
||||
value.clone()
|
||||
}
|
||||
|
||||
fn coerce_string_value(s: &str, schema: &serde_json::Value) -> Option<serde_json::Value> {
|
||||
if schema_allows_type(schema, "string") {
|
||||
return None;
|
||||
}
|
||||
|
||||
if schema_allows_type(schema, "integer")
|
||||
&& let Ok(v) = s.parse::<i64>()
|
||||
{
|
||||
return Some(serde_json::Value::from(v));
|
||||
}
|
||||
|
||||
if schema_allows_type(schema, "number")
|
||||
&& let Ok(v) = s.parse::<f64>()
|
||||
{
|
||||
return Some(serde_json::Value::from(v));
|
||||
}
|
||||
|
||||
if schema_allows_type(schema, "boolean") {
|
||||
match s.to_lowercase().as_str() {
|
||||
"true" => return Some(serde_json::json!(true)),
|
||||
"false" => return Some(serde_json::json!(false)),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if schema_allows_type(schema, "array") || schema_allows_type(schema, "object") {
|
||||
let parsed = serde_json::from_str::<serde_json::Value>(s).ok()?;
|
||||
let matches_schema = match &parsed {
|
||||
serde_json::Value::Array(_) => schema_allows_type(schema, "array"),
|
||||
serde_json::Value::Object(_) => schema_allows_type(schema, "object"),
|
||||
_ => false,
|
||||
};
|
||||
|
||||
if matches_schema {
|
||||
return Some(coerce_value(&parsed, schema));
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn schema_allows_type(schema: &serde_json::Value, expected: &str) -> bool {
|
||||
match schema.get("type") {
|
||||
Some(serde_json::Value::String(t)) => t == expected,
|
||||
Some(serde_json::Value::Array(types)) => types.iter().any(|t| t.as_str() == Some(expected)),
|
||||
_ => match expected {
|
||||
"object" => schema
|
||||
.get("properties")
|
||||
.and_then(|p| p.as_object())
|
||||
.is_some(),
|
||||
"array" => schema.get("items").is_some(),
|
||||
_ => false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::*;
|
||||
use crate::context::JobContext;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
|
||||
struct StubTool {
|
||||
schema: serde_json::Value,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for StubTool {
|
||||
fn name(&self) -> &str {
|
||||
"stub"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"stub"
|
||||
}
|
||||
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
self.schema.clone()
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
Ok(ToolOutput::success(params, Duration::from_millis(1)))
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn coerces_scalar_strings() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"count": { "type": "number" },
|
||||
"limit": { "type": "integer" },
|
||||
"enabled": { "type": "boolean" }
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({
|
||||
"count": "5",
|
||||
"limit": "10",
|
||||
"enabled": "TRUE"
|
||||
});
|
||||
|
||||
let result = prepare_params_for_schema(¶ms, &schema);
|
||||
|
||||
assert_eq!(result["count"], serde_json::json!(5.0)); // safety: test-only assertion
|
||||
assert_eq!(result["limit"], serde_json::json!(10)); // safety: test-only assertion
|
||||
assert_eq!(result["enabled"], serde_json::json!(true)); // safety: test-only assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn coerces_stringified_array_and_recurses_into_items() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"values": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": { "type": "integer" }
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({
|
||||
"values": "[[\"1\", \"2\"], [\"3\", 4]]"
|
||||
});
|
||||
|
||||
let result = prepare_params_for_schema(¶ms, &schema);
|
||||
|
||||
assert_eq!(result["values"], serde_json::json!([[1, 2], [3, 4]])); // safety: test-only assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn coerces_stringified_object_and_recurses_into_properties() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"request": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"start_index": { "type": "integer" },
|
||||
"enabled": { "type": ["boolean", "null"] }
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({
|
||||
"request": "{\"start_index\":\"12\",\"enabled\":\"false\"}"
|
||||
});
|
||||
|
||||
let result = prepare_params_for_schema(¶ms, &schema);
|
||||
|
||||
#[rustfmt::skip]
|
||||
assert_eq!( // safety: test-only assertion
|
||||
result["request"],
|
||||
serde_json::json!({"start_index": 12, "enabled": false})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn coerces_nullable_stringified_arrays() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"requests": {
|
||||
"type": ["array", "null"],
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"enabled": { "type": "boolean" }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({
|
||||
"requests": "[{\"enabled\":\"true\"}]"
|
||||
});
|
||||
|
||||
let result = prepare_params_for_schema(¶ms, &schema);
|
||||
|
||||
assert_eq!(result["requests"], serde_json::json!([{ "enabled": true }])); // safety: test-only assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn coerces_typed_additional_properties() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"additionalProperties": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"count": { "type": "integer" },
|
||||
"enabled": { "type": "boolean" }
|
||||
}
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({
|
||||
"alpha": "{\"count\":\"5\",\"enabled\":\"false\"}",
|
||||
"beta": { "count": "7", "enabled": "true" }
|
||||
});
|
||||
|
||||
let result = prepare_params_for_schema(¶ms, &schema);
|
||||
|
||||
#[rustfmt::skip]
|
||||
assert_eq!( // safety: test-only assertion
|
||||
result,
|
||||
serde_json::json!({
|
||||
"alpha": { "count": 5, "enabled": false },
|
||||
"beta": { "count": 7, "enabled": true }
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn leaves_invalid_json_strings_unchanged() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"requests": {
|
||||
"type": "array",
|
||||
"items": { "type": "object" }
|
||||
}
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({
|
||||
"requests": "[{\"oops\":]"
|
||||
});
|
||||
|
||||
let result = prepare_params_for_schema(¶ms, &schema);
|
||||
|
||||
assert_eq!(result["requests"], serde_json::json!("[{\"oops\":]")); // safety: test-only assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn leaves_string_when_schema_allows_string() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"value": { "type": ["string", "object"] }
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({
|
||||
"value": "{\"mode\":\"raw\"}"
|
||||
});
|
||||
|
||||
let result = prepare_params_for_schema(¶ms, &schema);
|
||||
|
||||
assert_eq!(result["value"], serde_json::json!("{\"mode\":\"raw\"}")); // safety: test-only assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn permissive_schema_is_noop() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"additionalProperties": true
|
||||
});
|
||||
let params = serde_json::json!({"count": "10"});
|
||||
|
||||
let result = prepare_params_for_schema(¶ms, &schema);
|
||||
|
||||
assert_eq!(result["count"], serde_json::json!("10")); // safety: test-only assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_tool_params_uses_discovery_schema() {
|
||||
let tool = StubTool {
|
||||
schema: serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"requests": {
|
||||
"type": "array",
|
||||
"items": { "type": "object" }
|
||||
}
|
||||
}
|
||||
}),
|
||||
};
|
||||
let params = serde_json::json!({
|
||||
"requests": "[{\"insertText\":{\"text\":\"hello\"}}]"
|
||||
});
|
||||
|
||||
let result = prepare_tool_params(&tool, ¶ms);
|
||||
|
||||
#[rustfmt::skip]
|
||||
assert_eq!( // safety: test-only assertion
|
||||
result["requests"],
|
||||
serde_json::json!([{ "insertText": { "text": "hello" } }])
|
||||
);
|
||||
}
|
||||
}
|
||||
+59
-4
@@ -8,7 +8,7 @@ use crate::context::JobContext;
|
||||
use crate::error::Error;
|
||||
use crate::llm::ChatMessage;
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::{ToolRegistry, redact_params};
|
||||
use crate::tools::{ToolRegistry, prepare_tool_params, redact_params};
|
||||
|
||||
/// Execute a tool with safety checks: lookup → validate → timeout → execute → serialize.
|
||||
///
|
||||
@@ -29,8 +29,10 @@ pub async fn execute_tool_with_safety(
|
||||
name: tool_name.to_string(),
|
||||
})?;
|
||||
|
||||
let normalized_params = prepare_tool_params(tool.as_ref(), params);
|
||||
|
||||
// Validate tool parameters
|
||||
let validation = safety.validator().validate_tool_params(params);
|
||||
let validation = safety.validator().validate_tool_params(&normalized_params);
|
||||
if !validation.is_valid {
|
||||
let details = validation
|
||||
.errors
|
||||
@@ -45,7 +47,7 @@ pub async fn execute_tool_with_safety(
|
||||
.into());
|
||||
}
|
||||
|
||||
let safe_params = redact_params(params, tool.sensitive_params());
|
||||
let safe_params = redact_params(&normalized_params, tool.sensitive_params());
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
params = %safe_params,
|
||||
@@ -56,7 +58,7 @@ pub async fn execute_tool_with_safety(
|
||||
let timeout = tool.execution_timeout();
|
||||
let start = std::time::Instant::now();
|
||||
let result = tokio::time::timeout(timeout, async {
|
||||
tool.execute(params.clone(), job_ctx).await
|
||||
tool.execute(normalized_params.clone(), job_ctx).await
|
||||
})
|
||||
.await;
|
||||
let elapsed = start.elapsed();
|
||||
@@ -237,6 +239,39 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
struct ArrayEchoTool;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for ArrayEchoTool {
|
||||
fn name(&self) -> &str {
|
||||
"array_echo"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"Echoes normalized params"
|
||||
}
|
||||
fn parameters_schema(&self) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"values": {
|
||||
"type": "array",
|
||||
"items": { "type": "integer" }
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
async fn execute(
|
||||
&self,
|
||||
params: serde_json::Value,
|
||||
_ctx: &JobContext,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
Ok(ToolOutput::success(params, Duration::default()))
|
||||
}
|
||||
fn requires_sanitization(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn test_safety() -> SafetyLayer {
|
||||
SafetyLayer::new(&crate::config::SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
@@ -348,6 +383,26 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_execute_normalizes_stringified_array_params() {
|
||||
let registry = registry_with(vec![Arc::new(ArrayEchoTool)]).await;
|
||||
let safety = test_safety();
|
||||
|
||||
let result = execute_tool_with_safety(
|
||||
®istry,
|
||||
&safety,
|
||||
"array_echo",
|
||||
&serde_json::json!({"values": "[\"1\", \"2\", 3]"}),
|
||||
&test_job_ctx(),
|
||||
)
|
||||
.await
|
||||
.expect("array_echo should succeed"); // safety: test-only assertion
|
||||
|
||||
let output: serde_json::Value =
|
||||
serde_json::from_str(&result).expect("tool result should be valid JSON"); // safety: test-only assertion
|
||||
assert_eq!(output["values"], serde_json::json!([1, 2, 3])); // safety: test-only assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_process_tool_result_success() {
|
||||
let safety = test_safety();
|
||||
|
||||
+27
-7
@@ -24,7 +24,7 @@ use crate::tools::mcp::config::McpServerConfig;
|
||||
/// Per-request timeouts can override the default via `.timeout()` on
|
||||
/// the request builder.
|
||||
fn oauth_http_client() -> Result<&'static reqwest::Client, AuthError> {
|
||||
static CLIENT: std::sync::OnceLock<Result<reqwest::Client, String>> =
|
||||
static CLIENT: std::sync::OnceLock<Result<reqwest::Client, AuthError>> =
|
||||
std::sync::OnceLock::new();
|
||||
CLIENT
|
||||
.get_or_init(|| {
|
||||
@@ -32,10 +32,10 @@ fn oauth_http_client() -> Result<&'static reqwest::Client, AuthError> {
|
||||
.timeout(Duration::from_secs(30))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|e| e.to_string())
|
||||
.map_err(|e| AuthError::Http(e.to_string()))
|
||||
})
|
||||
.as_ref()
|
||||
.map_err(|e| AuthError::Http(e.clone()))
|
||||
.map_err(Clone::clone)
|
||||
}
|
||||
|
||||
/// Log a debug message when a discovery/auth response is a redirect.
|
||||
@@ -57,7 +57,7 @@ fn log_redirect_if_applicable(url: &str, response: &reqwest::Response) {
|
||||
}
|
||||
|
||||
/// OAuth authorization error.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[derive(Debug, Clone, thiserror::Error)]
|
||||
pub enum AuthError {
|
||||
#[error("Server does not support OAuth authorization")]
|
||||
NotSupported,
|
||||
@@ -443,6 +443,11 @@ async fn fetch_resource_metadata(url: &str) -> Result<ProtectedResourceMetadata,
|
||||
}
|
||||
|
||||
/// Try to discover OAuth metadata via 401 challenge response.
|
||||
///
|
||||
/// Also accepts 400 responses, since some servers return 400 for
|
||||
/// unauthenticated requests. In practice the 400 path rarely yields a
|
||||
/// `WWW-Authenticate` header (GitHub's MCP does not), so discovery
|
||||
/// typically falls through to strategy 2 (RFC 9728) or 3 (direct).
|
||||
async fn discover_via_401(server_url: &str) -> Result<AuthorizationServerMetadata, AuthError> {
|
||||
validate_url_safe(server_url).await?;
|
||||
|
||||
@@ -459,9 +464,13 @@ async fn discover_via_401(server_url: &str) -> Result<AuthorizationServerMetadat
|
||||
|
||||
log_redirect_if_applicable(server_url, &response);
|
||||
|
||||
if response.status().as_u16() != 401 {
|
||||
let status = response.status().as_u16();
|
||||
|
||||
// Accept 401 (standard) and 400 (some servers like GitHub MCP use this).
|
||||
// In both cases, look for WWW-Authenticate header with discovery metadata.
|
||||
if status != 401 && status != 400 {
|
||||
return Err(AuthError::DiscoveryFailed(format!(
|
||||
"Expected 401, got {}",
|
||||
"Expected 401 or 400, got {}",
|
||||
response.status()
|
||||
)));
|
||||
}
|
||||
@@ -471,7 +480,7 @@ async fn discover_via_401(server_url: &str) -> Result<AuthorizationServerMetadat
|
||||
.get("WWW-Authenticate")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.ok_or_else(|| {
|
||||
AuthError::DiscoveryFailed("No WWW-Authenticate header in 401 response".to_string())
|
||||
AuthError::DiscoveryFailed(format!("No WWW-Authenticate header in {} response", status))
|
||||
})?;
|
||||
|
||||
let resource_metadata_url = parse_resource_metadata_url(www_auth).ok_or_else(|| {
|
||||
@@ -1511,6 +1520,17 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_auth_error_clone_preserves_http_variant_and_payload() {
|
||||
let original = AuthError::Http("builder failed".to_string());
|
||||
let cloned = original.clone();
|
||||
|
||||
match cloned {
|
||||
AuthError::Http(message) => assert_eq!(message, "builder failed"), // safety: test assertion in #[cfg(test)] module; not production panic path
|
||||
other => panic!("expected AuthError::Http variant, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
// --- New tests for well-known URI construction ---
|
||||
|
||||
#[test]
|
||||
|
||||
+142
-2
@@ -275,7 +275,10 @@ impl McpClient {
|
||||
.keys()
|
||||
.any(|k| k.eq_ignore_ascii_case("authorization"));
|
||||
if !has_custom_auth && let Some(token) = self.get_access_token().await? {
|
||||
headers.insert("Authorization".to_string(), format!("Bearer {}", token));
|
||||
let trimmed = token.trim();
|
||||
if !trimmed.is_empty() {
|
||||
headers.insert("Authorization".to_string(), format!("Bearer {}", trimmed));
|
||||
}
|
||||
}
|
||||
if let Some(ref session_manager) = self.session_manager
|
||||
&& let Some(session_id) = session_manager.get_session_id(&self.server_name).await
|
||||
@@ -302,7 +305,12 @@ impl McpClient {
|
||||
match result {
|
||||
Ok(response) => return Ok(response),
|
||||
Err(ToolError::ExternalService(ref msg))
|
||||
if msg.contains("401") || msg.contains("Unauthorized") =>
|
||||
if msg.contains("401")
|
||||
|| msg.contains("Unauthorized")
|
||||
|| (msg.contains("400") && {
|
||||
let lower = msg.to_ascii_lowercase();
|
||||
lower.contains("authorization") || lower.contains("authenticate")
|
||||
}) =>
|
||||
{
|
||||
if attempt == 0
|
||||
&& let Some(ref secrets) = self.secrets
|
||||
@@ -1113,4 +1121,136 @@ mod tests {
|
||||
let approval = wrapper.requires_approval(&serde_json::json!({}));
|
||||
assert_eq!(approval, ApprovalRequirement::Never);
|
||||
}
|
||||
|
||||
// Regression test: empty/whitespace-only tokens must not produce a
|
||||
// malformed `Authorization: Bearer ` header (GitHub MCP returns 400
|
||||
// "Authorization header is badly formatted" in this case).
|
||||
#[tokio::test]
|
||||
async fn test_build_headers_skips_empty_token() {
|
||||
use crate::secrets::{CreateSecretParams, DecryptedSecret, Secret, SecretError, SecretRef};
|
||||
use uuid::Uuid;
|
||||
|
||||
// In-memory secrets store that returns a whitespace-only string for the token.
|
||||
struct EmptyTokenStore;
|
||||
#[async_trait]
|
||||
impl crate::secrets::SecretsStore for EmptyTokenStore {
|
||||
async fn create(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_params: CreateSecretParams,
|
||||
) -> Result<Secret, SecretError> {
|
||||
unimplemented!()
|
||||
}
|
||||
async fn get(&self, _user_id: &str, _name: &str) -> Result<Secret, SecretError> {
|
||||
unimplemented!()
|
||||
}
|
||||
async fn get_decrypted(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_name: &str,
|
||||
) -> Result<DecryptedSecret, SecretError> {
|
||||
DecryptedSecret::from_bytes(b" ".to_vec())
|
||||
}
|
||||
async fn exists(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
|
||||
Ok(true)
|
||||
}
|
||||
async fn delete(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
|
||||
Ok(true)
|
||||
}
|
||||
async fn list(&self, _user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
async fn record_usage(&self, _secret_id: Uuid) -> Result<(), SecretError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn is_accessible(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_secret_name: &str,
|
||||
_allowed_secrets: &[String],
|
||||
) -> Result<bool, SecretError> {
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
let config = McpServerConfig::new("github", "https://api.githubcopilot.com/mcp/");
|
||||
let session_manager = Arc::new(McpSessionManager::new());
|
||||
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
|
||||
Arc::new(EmptyTokenStore);
|
||||
|
||||
let client = McpClient::new_authenticated(config, session_manager, secrets, "test-user");
|
||||
|
||||
let headers = client.build_request_headers().await.unwrap(); // safety: test
|
||||
assert!(
|
||||
// safety: test
|
||||
!headers.contains_key("Authorization"),
|
||||
"Empty/whitespace token must not produce an Authorization header, got: {:?}",
|
||||
headers.get("Authorization")
|
||||
);
|
||||
}
|
||||
|
||||
// Regression test: tokens with leading/trailing whitespace must be trimmed
|
||||
// before being used in the Authorization header.
|
||||
#[tokio::test]
|
||||
async fn test_build_headers_trims_token() {
|
||||
use crate::secrets::{CreateSecretParams, DecryptedSecret, Secret, SecretError, SecretRef};
|
||||
use uuid::Uuid;
|
||||
|
||||
struct PaddedTokenStore;
|
||||
#[async_trait]
|
||||
impl crate::secrets::SecretsStore for PaddedTokenStore {
|
||||
async fn create(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_params: CreateSecretParams,
|
||||
) -> Result<Secret, SecretError> {
|
||||
unimplemented!()
|
||||
}
|
||||
async fn get(&self, _user_id: &str, _name: &str) -> Result<Secret, SecretError> {
|
||||
unimplemented!()
|
||||
}
|
||||
async fn get_decrypted(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_name: &str,
|
||||
) -> Result<DecryptedSecret, SecretError> {
|
||||
DecryptedSecret::from_bytes(b" gho_abc123 \n".to_vec())
|
||||
}
|
||||
async fn exists(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
|
||||
Ok(true)
|
||||
}
|
||||
async fn delete(&self, _user_id: &str, _name: &str) -> Result<bool, SecretError> {
|
||||
Ok(true)
|
||||
}
|
||||
async fn list(&self, _user_id: &str) -> Result<Vec<SecretRef>, SecretError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
async fn record_usage(&self, _secret_id: Uuid) -> Result<(), SecretError> {
|
||||
Ok(())
|
||||
}
|
||||
async fn is_accessible(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
_secret_name: &str,
|
||||
_allowed_secrets: &[String],
|
||||
) -> Result<bool, SecretError> {
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
let config = McpServerConfig::new("github", "https://api.githubcopilot.com/mcp/");
|
||||
let session_manager = Arc::new(McpSessionManager::new());
|
||||
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
|
||||
Arc::new(PaddedTokenStore);
|
||||
|
||||
let client = McpClient::new_authenticated(config, session_manager, secrets, "test-user");
|
||||
|
||||
let headers = client.build_request_headers().await.unwrap(); // safety: test
|
||||
assert_eq!(
|
||||
// safety: test
|
||||
headers.get("Authorization").unwrap(), // safety: test
|
||||
"Bearer gho_abc123",
|
||||
"Token must be trimmed before use in Authorization header"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,7 +39,7 @@ impl HttpMcpTransport {
|
||||
http_client: reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.build()
|
||||
.expect("Failed to create HTTP client"),
|
||||
.expect("Failed to create HTTP client"), // safety: TLS init with default rustls cannot fail
|
||||
session_manager: None,
|
||||
custom_headers: HashMap::new(),
|
||||
}
|
||||
@@ -212,9 +212,10 @@ impl HttpMcpTransport {
|
||||
}
|
||||
}
|
||||
}
|
||||
// Keep only the unprocessed trailing fragment.
|
||||
// Keep only the unprocessed trailing fragment without allocating
|
||||
// a new String each iteration.
|
||||
if remaining_start > 0 {
|
||||
buffer = buffer[remaining_start..].to_string();
|
||||
buffer.drain(..remaining_start);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
|
||||
pub mod builder;
|
||||
pub mod builtin;
|
||||
mod coercion;
|
||||
pub mod execute;
|
||||
pub mod mcp;
|
||||
pub mod rate_limiter;
|
||||
@@ -24,6 +25,7 @@ pub use builder::{
|
||||
LlmSoftwareBuilder, SoftwareBuilder, SoftwareType, Template, TemplateEngine, TemplateType,
|
||||
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
|
||||
};
|
||||
pub(crate) use coercion::prepare_tool_params;
|
||||
pub use rate_limiter::RateLimiter;
|
||||
pub use registry::ToolRegistry;
|
||||
pub use tool::{
|
||||
|
||||
+325
-170
@@ -343,7 +343,7 @@ impl near::agent::host::Host for StoreData {
|
||||
.map_err(|e| format!("Failed to create HTTP runtime: {e}"))?,
|
||||
);
|
||||
}
|
||||
let rt = self.http_runtime.as_ref().expect("just initialized");
|
||||
let rt = self.http_runtime.as_ref().expect("just initialized"); // safety: is_none branch above guarantees Some
|
||||
let result = rt.block_on(async {
|
||||
let client = reqwest::Client::builder()
|
||||
.connect_timeout(Duration::from_secs(10))
|
||||
@@ -485,7 +485,7 @@ struct WasmToolSchemas {
|
||||
/// This stays permissive by default to avoid serializing full exported
|
||||
/// WASM schemas on every LLM call. Sidecars can override it explicitly.
|
||||
advertised: serde_json::Value,
|
||||
/// Full schema available for discovery and coercion.
|
||||
/// Full schema available for discovery and runtime parameter preparation.
|
||||
///
|
||||
/// Seeded from the WASM `schema()` export at registration time, unless a
|
||||
/// sidecar explicitly overrides it.
|
||||
@@ -508,6 +508,19 @@ impl WasmToolSchemas {
|
||||
.is_none_or(|p| p.is_empty())
|
||||
}
|
||||
|
||||
fn typed_property_count(schema: &serde_json::Value) -> usize {
|
||||
schema
|
||||
.get("properties")
|
||||
.and_then(|p| p.as_object())
|
||||
.map(|props| {
|
||||
props
|
||||
.values()
|
||||
.filter(|prop| schema_is_typed_property(prop))
|
||||
.count()
|
||||
})
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn new(discovery: serde_json::Value) -> Self {
|
||||
Self {
|
||||
advertised: Self::permissive_schema(),
|
||||
@@ -533,27 +546,6 @@ impl WasmToolSchemas {
|
||||
fn discovery(&self) -> serde_json::Value {
|
||||
self.discovery.clone()
|
||||
}
|
||||
|
||||
/// Return the best schema available for type coercion.
|
||||
///
|
||||
/// Prefers the discovery schema when it has typed properties. Falls back
|
||||
/// to the `PreparedModule` schema extracted at load time rather than
|
||||
/// re-calling the WASM `schema()` export mid-execution, which could
|
||||
/// interact with mutable linear memory state.
|
||||
fn effective_for_coercion(&self, prepared_schema: &serde_json::Value) -> serde_json::Value {
|
||||
if !Self::is_permissive_schema(&self.discovery) {
|
||||
return self.discovery.clone();
|
||||
}
|
||||
|
||||
// Fall back to the load-time extracted schema from PreparedModule.
|
||||
// This avoids calling schema() on the already-running WASM instance
|
||||
// where mutable state could produce inconsistent results.
|
||||
if !Self::is_permissive_schema(prepared_schema) {
|
||||
return prepared_schema.clone();
|
||||
}
|
||||
|
||||
self.discovery.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl WasmToolWrapper {
|
||||
@@ -583,7 +575,21 @@ impl WasmToolWrapper {
|
||||
|
||||
/// Override the parameter schema.
|
||||
pub fn with_schema(mut self, schema: serde_json::Value) -> Self {
|
||||
self.schemas = self.schemas.with_override(schema);
|
||||
let override_typed = WasmToolSchemas::typed_property_count(&schema);
|
||||
let prepared_typed = WasmToolSchemas::typed_property_count(&self.prepared.schema);
|
||||
|
||||
if override_typed == 0 && prepared_typed > 0 {
|
||||
tracing::warn!(
|
||||
tool = %self.prepared.name,
|
||||
"Ignoring untyped schema override for discovery/runtime preparation and preserving extracted WASM schema"
|
||||
);
|
||||
self.schemas = WasmToolSchemas {
|
||||
advertised: schema,
|
||||
discovery: self.prepared.schema.clone(),
|
||||
};
|
||||
} else {
|
||||
self.schemas = self.schemas.with_override(schema);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
@@ -697,16 +703,6 @@ impl WasmToolWrapper {
|
||||
// Get typed interface — used for execute.
|
||||
let tool_iface = instance.near_agent_tool();
|
||||
|
||||
// Determine effective schema for type coercion.
|
||||
// Prefer the discovery schema when typed; fall back to the load-time
|
||||
// extracted schema from PreparedModule rather than re-calling the WASM
|
||||
// export on the already-running instance.
|
||||
let effective_schema = self.schemas.effective_for_coercion(&self.prepared.schema);
|
||||
|
||||
// Coerce string-encoded values to their schema-declared types.
|
||||
// LLMs frequently pass numeric values as strings (e.g. "5" instead of 5).
|
||||
let params = coerce_params_to_schema(params, &effective_schema);
|
||||
|
||||
// Prepare the request
|
||||
let params_json = serde_json::to_string(¶ms)
|
||||
.map_err(|e| WasmError::InvalidResponseJson(e.to_string()))?;
|
||||
@@ -734,10 +730,7 @@ impl WasmToolWrapper {
|
||||
// Check for tool-level error — point the LLM to tool_info for the
|
||||
// full schema instead of dumping ~3.5KB inline.
|
||||
if let Some(err) = response.error {
|
||||
let hint = format!(
|
||||
"Tip: call tool_info(name: \"{}\", include_schema: true) for the full parameter schema.",
|
||||
self.prepared.name
|
||||
);
|
||||
let hint = build_tool_usage_hint(&self.prepared.name, &self.schemas.discovery());
|
||||
return Err(WasmError::ToolReturnedError { message: err, hint });
|
||||
}
|
||||
|
||||
@@ -1104,7 +1097,18 @@ async fn resolve_host_credentials(
|
||||
) -> Vec<ResolvedHostCredential> {
|
||||
let store = match store {
|
||||
Some(s) => s,
|
||||
None => return Vec::new(),
|
||||
None => {
|
||||
// If tool requires credentials but has no secrets store, this is a configuration error
|
||||
if let Some(http_cap) = &capabilities.http
|
||||
&& !http_cap.credentials.is_empty()
|
||||
{
|
||||
tracing::warn!(
|
||||
user_id = %user_id,
|
||||
"WASM tool requires credentials but secrets_store is not configured - authentication will fail"
|
||||
);
|
||||
}
|
||||
return Vec::new();
|
||||
}
|
||||
};
|
||||
|
||||
// Check if the access token needs refreshing before resolving credentials.
|
||||
@@ -1155,13 +1159,37 @@ async fn resolve_host_credentials(
|
||||
continue;
|
||||
}
|
||||
|
||||
// Try to get credential under the provided user_id first.
|
||||
// If not found and user_id != "default", fallback to "default" (global credentials).
|
||||
// This handles OAuth tokens stored globally under "default" but accessed from routine contexts.
|
||||
let secret = match store.get_decrypted(user_id, &mapping.secret_name).await {
|
||||
Ok(s) => s,
|
||||
Ok(s) => Some(s),
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
// If lookup fails and we're not already looking up "default", try "default" as fallback
|
||||
if user_id != "default" {
|
||||
tracing::debug!(
|
||||
secret_name = %mapping.secret_name,
|
||||
user_id = %user_id,
|
||||
error = %e,
|
||||
"Credential not found for user, trying default global credentials"
|
||||
);
|
||||
store
|
||||
.get_decrypted("default", &mapping.secret_name)
|
||||
.await
|
||||
.ok()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let secret = match secret {
|
||||
Some(s) => s,
|
||||
None => {
|
||||
tracing::warn!(
|
||||
secret_name = %mapping.secret_name,
|
||||
error = %e,
|
||||
"Could not resolve credential for WASM tool (auth may not be configured)"
|
||||
user_id = %user_id,
|
||||
"Could not resolve credential for WASM tool (not found in user context or default)"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
@@ -1290,59 +1318,69 @@ fn is_private_ip(ip: std::net::IpAddr) -> bool {
|
||||
}
|
||||
}
|
||||
|
||||
/// Coerce parameter values to match their JSON Schema-declared types.
|
||||
///
|
||||
/// LLMs frequently send numeric values as strings (e.g. `"5"` instead of `5`)
|
||||
/// or booleans as strings (`"true"` instead of `true`). This walks the params
|
||||
/// object and converts string values where the schema expects a different type.
|
||||
fn coerce_params_to_schema(
|
||||
mut params: serde_json::Value,
|
||||
schema: &serde_json::Value,
|
||||
) -> serde_json::Value {
|
||||
let properties = schema.get("properties").and_then(|p| p.as_object());
|
||||
fn schema_contains_container_properties(schema: &serde_json::Value) -> bool {
|
||||
schema
|
||||
.get("properties")
|
||||
.and_then(|p| p.as_object())
|
||||
.map(|props| {
|
||||
props.values().any(|prop| {
|
||||
schema_declares_type(prop, "array") || schema_declares_type(prop, "object")
|
||||
})
|
||||
})
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
let properties = match properties {
|
||||
Some(p) => p,
|
||||
None => return params,
|
||||
};
|
||||
|
||||
let obj = match params.as_object_mut() {
|
||||
Some(o) => o,
|
||||
None => return params,
|
||||
};
|
||||
|
||||
for (key, prop_schema) in properties {
|
||||
let declared_type = prop_schema.get("type").and_then(|t| t.as_str());
|
||||
let declared_type = match declared_type {
|
||||
Some(t) => t,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
if let Some(current_value) = obj.get_mut(key)
|
||||
&& let Some(s) = current_value.as_str()
|
||||
{
|
||||
if declared_type == "string" {
|
||||
continue;
|
||||
fn schema_declares_type(schema: &serde_json::Value, expected: &str) -> bool {
|
||||
match schema.get("type") {
|
||||
Some(serde_json::Value::String(t)) => t == expected,
|
||||
Some(serde_json::Value::Array(types)) => types.iter().any(|t| t.as_str() == Some(expected)),
|
||||
_ => match expected {
|
||||
"object" => {
|
||||
schema
|
||||
.get("properties")
|
||||
.and_then(|p| p.as_object())
|
||||
.is_some()
|
||||
|| schema
|
||||
.get("additionalProperties")
|
||||
.is_some_and(serde_json::Value::is_object)
|
||||
}
|
||||
"array" => schema.get("items").is_some(),
|
||||
_ => false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
let coerced = match declared_type {
|
||||
"number" => s.parse::<f64>().ok().map(serde_json::Value::from),
|
||||
"integer" => s.parse::<i64>().ok().map(serde_json::Value::from),
|
||||
"boolean" => match s.to_lowercase().as_str() {
|
||||
"true" => Some(serde_json::json!(true)),
|
||||
"false" => Some(serde_json::json!(false)),
|
||||
_ => None,
|
||||
},
|
||||
_ => None,
|
||||
};
|
||||
fn schema_is_typed_property(schema: &serde_json::Value) -> bool {
|
||||
matches!(
|
||||
schema.get("type"),
|
||||
Some(serde_json::Value::String(_)) | Some(serde_json::Value::Array(_))
|
||||
) || schema.get("$ref").is_some()
|
||||
|| schema.get("anyOf").is_some()
|
||||
|| schema.get("oneOf").is_some()
|
||||
|| schema.get("allOf").is_some()
|
||||
|| schema.get("items").is_some()
|
||||
|| schema
|
||||
.get("properties")
|
||||
.and_then(|p| p.as_object())
|
||||
.is_some()
|
||||
|| schema
|
||||
.get("additionalProperties")
|
||||
.is_some_and(serde_json::Value::is_object)
|
||||
}
|
||||
|
||||
if let Some(new_val) = coerced {
|
||||
*current_value = new_val;
|
||||
}
|
||||
}
|
||||
fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String {
|
||||
let mut hint = format!(
|
||||
"Tip: call tool_info(name: \"{}\", include_schema: true) for the full parameter schema.",
|
||||
tool_name
|
||||
);
|
||||
|
||||
if schema_contains_container_properties(schema) {
|
||||
hint.push_str(
|
||||
" For array/object fields, pass native JSON arrays/objects, not quoted JSON strings.",
|
||||
);
|
||||
}
|
||||
|
||||
params
|
||||
hint
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -1910,100 +1948,60 @@ mod tests {
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_params_string_to_number() {
|
||||
let schema = serde_json::json!({
|
||||
#[tokio::test]
|
||||
async fn test_untyped_override_preserves_extracted_discovery_schema() {
|
||||
let typed_schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"count": { "type": "number" },
|
||||
"name": { "type": "string" }
|
||||
"values": {
|
||||
"type": ["array", "null"],
|
||||
"items": { "type": "array" }
|
||||
}
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({"count": "5", "name": "test"});
|
||||
let result = super::coerce_params_to_schema(params, &schema);
|
||||
assert_eq!(result["count"], serde_json::json!(5.0));
|
||||
assert_eq!(result["name"], serde_json::json!("test"));
|
||||
|
||||
let runtime = Arc::new(WasmToolRuntime::new(WasmRuntimeConfig::for_testing()).unwrap()); // safety: test-only setup
|
||||
let mut prepared = runtime
|
||||
.prepare("sheets", b"\0asm\x0d\0\x01\0", None)
|
||||
.await
|
||||
.unwrap(); // safety: test-only setup
|
||||
Arc::get_mut(&mut prepared).unwrap().schema = typed_schema.clone(); // safety: test-only setup
|
||||
|
||||
let wrapper =
|
||||
super::WasmToolWrapper::new(Arc::clone(&runtime), prepared, Capabilities::default())
|
||||
.with_schema(serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"additionalProperties": true
|
||||
}));
|
||||
|
||||
#[rustfmt::skip]
|
||||
assert_eq!( // safety: test-only assertion
|
||||
wrapper.parameters_schema(),
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"additionalProperties": true
|
||||
})
|
||||
);
|
||||
assert_eq!(wrapper.discovery_schema(), typed_schema); // safety: test-only assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_params_string_to_integer() {
|
||||
fn test_build_tool_usage_hint_detects_nullable_container_properties() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"limit": { "type": "integer" }
|
||||
"requests": {
|
||||
"type": ["array", "null"],
|
||||
"items": { "type": "object" }
|
||||
}
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({"limit": "10"});
|
||||
let result = super::coerce_params_to_schema(params, &schema);
|
||||
assert_eq!(result["limit"], serde_json::json!(10));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_params_string_to_boolean() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": { "type": "boolean" },
|
||||
"b": { "type": "boolean" },
|
||||
"c": { "type": "boolean" },
|
||||
"d": { "type": "boolean" }
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({
|
||||
"a": "true",
|
||||
"b": "false",
|
||||
"c": "True",
|
||||
"d": "FALSE"
|
||||
});
|
||||
let result = super::coerce_params_to_schema(params, &schema);
|
||||
assert_eq!(result["a"], serde_json::json!(true));
|
||||
assert_eq!(result["b"], serde_json::json!(false));
|
||||
assert_eq!(result["c"], serde_json::json!(true));
|
||||
assert_eq!(result["d"], serde_json::json!(false));
|
||||
}
|
||||
let hint = super::build_tool_usage_hint("google_docs", &schema);
|
||||
|
||||
#[test]
|
||||
fn test_coerce_params_already_correct_type() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"count": { "type": "number" }
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({"count": 5});
|
||||
let result = super::coerce_params_to_schema(params, &schema);
|
||||
assert_eq!(result["count"], serde_json::json!(5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_params_invalid_string_not_coerced() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"count": { "type": "number" }
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({"count": "not-a-number"});
|
||||
let result = super::coerce_params_to_schema(params, &schema);
|
||||
// Should remain as string since it can't be parsed
|
||||
assert_eq!(result["count"], serde_json::json!("not-a-number"));
|
||||
}
|
||||
|
||||
/// Regression: permissive fallback schema (empty properties) must NOT coerce.
|
||||
/// This documents the bug where WASM tools with no sidecar `parameters` field
|
||||
/// got the permissive fallback, causing coercion to be a no-op and LLM-provided
|
||||
/// string integers to reach the WASM tool un-coerced.
|
||||
#[test]
|
||||
fn test_coerce_noop_with_permissive_schema() {
|
||||
let permissive = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"additionalProperties": true
|
||||
});
|
||||
let params = serde_json::json!({"query": "test", "count": "10"});
|
||||
let result = super::coerce_params_to_schema(params, &permissive);
|
||||
// With empty properties, no coercion happens — string stays string
|
||||
assert_eq!(result["count"], serde_json::json!("10"));
|
||||
assert!(hint.contains("native JSON arrays/objects")); // safety: test-only assertion
|
||||
}
|
||||
|
||||
/// Regression test: leak scan must run on raw headers (before credential
|
||||
@@ -2058,4 +2056,161 @@ mod tests {
|
||||
"Leak scan on post-injection headers should block the Slack token"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_fallback_to_default_user() {
|
||||
use crate::secrets::{CredentialLocation, CredentialMapping, SecretsStore};
|
||||
use crate::tools::wasm::capabilities::HttpCapability;
|
||||
use crate::tools::wasm::wrapper::resolve_host_credentials;
|
||||
|
||||
let store = test_secrets_store();
|
||||
|
||||
// Store a token under the "default" global user
|
||||
store
|
||||
.create(
|
||||
"default",
|
||||
crate::secrets::CreateSecretParams::new("google_oauth_token", "global_token_value"),
|
||||
)
|
||||
.await
|
||||
.expect("Failed to store global token"); // safety: test code only
|
||||
|
||||
// Create capabilities requiring this credential
|
||||
let mut creds = std::collections::HashMap::new();
|
||||
creds.insert(
|
||||
"google_oauth_token".to_string(),
|
||||
CredentialMapping {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
host_patterns: vec!["sheets.googleapis.com".to_string()],
|
||||
},
|
||||
);
|
||||
let caps = Capabilities {
|
||||
http: Some(HttpCapability {
|
||||
allowlist: vec![],
|
||||
credentials: creds,
|
||||
rate_limit: crate::tools::wasm::capabilities::RateLimitConfig::default(),
|
||||
max_request_bytes: 1024 * 1024,
|
||||
max_response_bytes: 10 * 1024 * 1024,
|
||||
timeout: std::time::Duration::from_secs(30),
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Resolve credentials for a different user (routine context)
|
||||
// Should fallback to "default" and find the token
|
||||
let result = resolve_host_credentials(&caps, Some(&store), "routine_user_123", None).await;
|
||||
|
||||
assert!(!result.is_empty(), "fallback to default"); // safety: test code only
|
||||
assert_eq!(result[0].secret_value, "global_token_value"); // safety: test code only
|
||||
}
|
||||
|
||||
fn test_capabilities_with_google_oauth() -> Capabilities {
|
||||
use crate::secrets::{CredentialLocation, CredentialMapping};
|
||||
use crate::tools::wasm::capabilities::HttpCapability;
|
||||
|
||||
let mut creds = std::collections::HashMap::new();
|
||||
creds.insert(
|
||||
"google_oauth_token".to_string(),
|
||||
CredentialMapping {
|
||||
secret_name: "google_oauth_token".to_string(),
|
||||
location: CredentialLocation::AuthorizationBearer,
|
||||
host_patterns: vec!["sheets.googleapis.com".to_string()],
|
||||
},
|
||||
);
|
||||
Capabilities {
|
||||
http: Some(HttpCapability {
|
||||
allowlist: vec![],
|
||||
credentials: creds,
|
||||
rate_limit: crate::tools::wasm::capabilities::RateLimitConfig::default(),
|
||||
max_request_bytes: 1024 * 1024,
|
||||
max_response_bytes: 10 * 1024 * 1024,
|
||||
timeout: std::time::Duration::from_secs(30),
|
||||
}),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_prefers_user_specific_over_default() {
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::tools::wasm::wrapper::resolve_host_credentials;
|
||||
|
||||
let store = test_secrets_store();
|
||||
|
||||
// Store token under "default" (global)
|
||||
store
|
||||
.create(
|
||||
"default",
|
||||
crate::secrets::CreateSecretParams::new("google_oauth_token", "global_token"),
|
||||
)
|
||||
.await
|
||||
.expect("Failed to store global token"); // safety: test code only
|
||||
|
||||
// Store token under user_123 (user-specific)
|
||||
store
|
||||
.create(
|
||||
"user_123",
|
||||
crate::secrets::CreateSecretParams::new(
|
||||
"google_oauth_token",
|
||||
"user_specific_token",
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("Failed to store user token"); // safety: test code only
|
||||
|
||||
// Create capabilities
|
||||
let caps = test_capabilities_with_google_oauth();
|
||||
|
||||
// Resolve credentials for user_123
|
||||
// Should prefer user_123's token over default
|
||||
let result = resolve_host_credentials(&caps, Some(&store), "user_123", None).await;
|
||||
|
||||
assert!(!result.is_empty(), "has user credentials"); // safety: test code only
|
||||
assert_eq!(result[0].secret_value, "user_specific_token", "user token"); // safety: test code only
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_no_fallback_when_already_default() {
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::tools::wasm::wrapper::resolve_host_credentials;
|
||||
|
||||
let store = test_secrets_store();
|
||||
|
||||
// Only store token under "default" (not a duplicate)
|
||||
store
|
||||
.create(
|
||||
"default",
|
||||
crate::secrets::CreateSecretParams::new("google_oauth_token", "default_token"),
|
||||
)
|
||||
.await
|
||||
.expect("Failed to store default token"); // safety: test code only
|
||||
|
||||
// Create capabilities
|
||||
let caps = test_capabilities_with_google_oauth();
|
||||
|
||||
// Resolve credentials for "default" user
|
||||
// Should NOT attempt fallback (already looking up default)
|
||||
let result = resolve_host_credentials(&caps, Some(&store), "default", None).await;
|
||||
|
||||
assert!(!result.is_empty(), "Should find default token"); // safety: test code only
|
||||
assert_eq!(result[0].secret_value, "default_token"); // safety: test code only
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resolve_host_credentials_missing_secret_warns() {
|
||||
use crate::tools::wasm::wrapper::resolve_host_credentials;
|
||||
|
||||
let store = test_secrets_store();
|
||||
|
||||
// Don't store any token
|
||||
|
||||
// Create capabilities expecting a credential
|
||||
let caps = test_capabilities_with_google_oauth();
|
||||
|
||||
// Resolve credentials when neither user nor default has the token
|
||||
let result = resolve_host_credentials(&caps, Some(&store), "user_456", None).await;
|
||||
|
||||
// Should return empty since credential can't be found anywhere
|
||||
assert!(result.is_empty(), "no credentials found"); // safety: test code only
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
};
|
||||
|
||||
|
||||
+27
-16
@@ -30,7 +30,7 @@ use crate::llm::{
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::execute::process_tool_result;
|
||||
use crate::tools::rate_limiter::RateLimitResult;
|
||||
use crate::tools::{ApprovalContext, ToolRegistry, redact_params};
|
||||
use crate::tools::{ApprovalContext, ToolRegistry, prepare_tool_params, redact_params};
|
||||
|
||||
/// Shared dependencies for worker execution.
|
||||
///
|
||||
@@ -483,8 +483,10 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
name: tool_name.to_string(),
|
||||
})?;
|
||||
|
||||
let normalized_params = prepare_tool_params(tool.as_ref(), params);
|
||||
|
||||
// Check approval: use context-aware check if available, else block all non-Never tools
|
||||
let requirement = tool.requires_approval(params);
|
||||
let requirement = tool.requires_approval(&normalized_params);
|
||||
let blocked =
|
||||
ApprovalContext::is_blocked_or_default(&deps.approval_context, tool_name, requirement);
|
||||
if blocked {
|
||||
@@ -517,9 +519,9 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
}
|
||||
|
||||
// Run BeforeToolCall hook
|
||||
let params = {
|
||||
let effective_params = {
|
||||
use crate::hooks::{HookError, HookEvent, HookOutcome};
|
||||
let hook_params = redact_params(params, tool.sensitive_params());
|
||||
let hook_params = redact_params(&normalized_params, tool.sensitive_params());
|
||||
let event = HookEvent::ToolCall {
|
||||
tool_name: tool_name.to_string(),
|
||||
parameters: hook_params,
|
||||
@@ -543,15 +545,21 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
}
|
||||
Ok(HookOutcome::Continue {
|
||||
modified: Some(new_params),
|
||||
}) => serde_json::from_str(&new_params).unwrap_or_else(|e| {
|
||||
tracing::warn!(
|
||||
tool = %tool_name,
|
||||
"Hook returned non-JSON modification for ToolCall, ignoring: {}",
|
||||
e
|
||||
);
|
||||
params.clone()
|
||||
}),
|
||||
_ => params.clone(),
|
||||
}) => match serde_json::from_str(&new_params) {
|
||||
// Hook output is fresh JSON text and may reintroduce stringified scalars or
|
||||
// containers, so we normalize it again. The fallback path reuses the already
|
||||
// normalized input because no hook mutation was applied.
|
||||
Ok(parsed) => prepare_tool_params(tool.as_ref(), &parsed),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
tool = %tool_name,
|
||||
"Hook returned non-JSON modification for ToolCall, ignoring: {}",
|
||||
e
|
||||
);
|
||||
normalized_params
|
||||
}
|
||||
},
|
||||
_ => normalized_params,
|
||||
}
|
||||
};
|
||||
if job_ctx.state == JobState::Cancelled {
|
||||
@@ -563,7 +571,10 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
}
|
||||
|
||||
// Validate tool parameters
|
||||
let validation = deps.safety.validator().validate_tool_params(¶ms);
|
||||
let validation = deps
|
||||
.safety
|
||||
.validator()
|
||||
.validate_tool_params(&effective_params);
|
||||
if !validation.is_valid {
|
||||
let details = validation
|
||||
.errors
|
||||
@@ -579,7 +590,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
}
|
||||
|
||||
// Redact sensitive parameter values before they touch any observability or audit path.
|
||||
let safe_params = redact_params(¶ms, tool.sensitive_params());
|
||||
let safe_params = redact_params(&effective_params, tool.sensitive_params());
|
||||
tracing::debug!(
|
||||
tool = %tool_name,
|
||||
params = %safe_params,
|
||||
@@ -591,7 +602,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
||||
let tool_timeout = tool.execution_timeout();
|
||||
let start = std::time::Instant::now();
|
||||
let result = tokio::time::timeout(tool_timeout, async {
|
||||
tool.execute(params.clone(), &job_ctx).await
|
||||
tool.execute(effective_params.clone(), &job_ctx).await
|
||||
})
|
||||
.await;
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
@@ -92,8 +92,9 @@ pub fn chunk_document(content: &str, config: ChunkConfig) -> Vec<String> {
|
||||
let chunk_words = &words[start..end];
|
||||
|
||||
// Don't create tiny trailing chunks, merge with previous
|
||||
if chunk_words.len() < config.min_chunk_size && !chunks.is_empty() {
|
||||
let last = chunks.pop().unwrap();
|
||||
if chunk_words.len() < config.min_chunk_size
|
||||
&& let Some(last) = chunks.pop()
|
||||
{
|
||||
let combined = format!("{} {}", last, chunk_words.join(" "));
|
||||
chunks.push(combined);
|
||||
break;
|
||||
|
||||
Reference in New Issue
Block a user