Merge branch 'staging' into feat/custom-llm-provider

This commit is contained in:
jinxin
2026-03-26 15:51:52 +08:00
committed by GitHub
65 changed files with 3313 additions and 272 deletions
+19 -5
View File
@@ -12,6 +12,7 @@ jobs:
tests:
name: Tests (${{ matrix.name }})
runs-on: ubuntu-latest
timeout-minutes: 45
strategy:
fail-fast: false
matrix:
@@ -40,11 +41,14 @@ jobs:
- name: Build WASM channels (for integration tests)
run: ./scripts/build-wasm-extensions.sh --channels
- name: Run Tests
run: cargo test ${{ matrix.flags }} -- --nocapture
run: |
timeout --signal=INT --kill-after=30s 40m \
cargo test ${{ matrix.flags }} -- --nocapture
heavy-integration-tests:
name: Heavy Integration Tests
runs-on: ubuntu-latest
timeout-minutes: 20
steps:
- name: Checkout repository
uses: actions/checkout@v6
@@ -58,9 +62,13 @@ jobs:
- name: Build Telegram WASM channel
run: cargo build --manifest-path channels-src/telegram/Cargo.toml --target wasm32-wasip2 --release
- name: Run thread scheduling integration tests
run: cargo test --no-default-features --features libsql,integration --test e2e_thread_scheduling -- --nocapture
run: |
timeout --signal=INT --kill-after=30s 15m \
cargo test --no-default-features --features libsql,integration --test e2e_thread_scheduling -- --nocapture
- name: Run Telegram thread-scope regression test
run: cargo test --features integration --test telegram_auth_integration test_private_messages_use_chat_id_as_thread_scope -- --exact
run: |
timeout --signal=INT --kill-after=30s 10m \
cargo test --features integration --test telegram_auth_integration test_private_messages_use_chat_id_as_thread_scope -- --exact
telegram-tests:
name: Telegram Channel Tests
@@ -68,6 +76,7 @@ jobs:
github.event_name != 'pull_request' ||
github.base_ref != 'staging'
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- name: Checkout repository
uses: actions/checkout@v6
@@ -75,7 +84,9 @@ jobs:
uses: dtolnay/rust-toolchain@stable
- uses: Swatinem/rust-cache@v2
- name: Run Telegram Channel Tests
run: cargo test --manifest-path channels-src/telegram/Cargo.toml -- --nocapture
run: |
timeout --signal=INT --kill-after=30s 10m \
cargo test --manifest-path channels-src/telegram/Cargo.toml -- --nocapture
windows-build:
name: Windows Build (${{ matrix.name }})
@@ -110,6 +121,7 @@ jobs:
github.event_name != 'pull_request' ||
github.base_ref != 'staging'
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- name: Checkout repository
uses: actions/checkout@v6
@@ -125,7 +137,9 @@ jobs:
- name: Build all WASM extensions against current WIT
run: ./scripts/build-wasm-extensions.sh
- name: Instantiation test (host linker compatibility)
run: cargo test --all-features wit_compat -- --nocapture
run: |
timeout --signal=INT --kill-after=30s 20m \
cargo test --all-features wit_compat -- --nocapture
bench-compile:
name: Benchmark Compilation
+55
View File
@@ -7,6 +7,32 @@
use serde::{Deserialize, Serialize};
/// A single tool decision in a reasoning update (SSE DTO).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDecisionDto {
pub tool_name: String,
pub rationale: String,
}
impl ToolDecisionDto {
/// Parse a list of tool decisions from a JSON array value.
pub fn from_json_array(value: &serde_json::Value) -> Vec<Self> {
value
.as_array()
.map(|arr| {
arr.iter()
.filter_map(|d| {
Some(Self {
tool_name: d.get("tool_name")?.as_str()?.to_string(),
rationale: d.get("rationale")?.as_str()?.to_string(),
})
})
.collect()
})
.unwrap_or_default()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum AppEvent {
@@ -163,6 +189,23 @@ pub enum AppEvent {
#[serde(skip_serializing_if = "Option::is_none")]
message: Option<String>,
},
/// Agent reasoning update (why it chose specific tools).
#[serde(rename = "reasoning_update")]
ReasoningUpdate {
narrative: String,
decisions: Vec<ToolDecisionDto>,
#[serde(skip_serializing_if = "Option::is_none")]
thread_id: Option<String>,
},
/// Reasoning update for a sandbox job.
#[serde(rename = "job_reasoning")]
JobReasoning {
job_id: String,
narrative: String,
decisions: Vec<ToolDecisionDto>,
},
}
impl AppEvent {
@@ -191,6 +234,8 @@ impl AppEvent {
Self::Suggestions { .. } => "suggestions",
Self::TurnCost { .. } => "turn_cost",
Self::ExtensionStatus { .. } => "extension_status",
Self::ReasoningUpdate { .. } => "reasoning_update",
Self::JobReasoning { .. } => "job_reasoning",
}
}
}
@@ -311,6 +356,16 @@ mod tests {
status: String::new(),
message: None,
},
AppEvent::ReasoningUpdate {
narrative: String::new(),
decisions: vec![],
thread_id: None,
},
AppEvent::JobReasoning {
job_id: String::new(),
narrative: String::new(),
decisions: vec![],
},
];
for variant in &variants {
+1 -1
View File
@@ -3,5 +3,5 @@
mod event;
mod util;
pub use event::AppEvent;
pub use event::{AppEvent, ToolDecisionDto};
pub use util::truncate_preview;
+215 -41
View File
@@ -13,9 +13,10 @@ use futures::StreamExt;
use uuid::Uuid;
use crate::agent::context_monitor::ContextMonitor;
use crate::agent::heartbeat::spawn_heartbeat;
use crate::agent::heartbeat::{spawn_heartbeat, spawn_multi_user_heartbeat};
use crate::agent::routine_engine::{RoutineEngine, spawn_cron_ticker};
use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair};
use crate::agent::session::ThreadState;
use crate::agent::session_manager::SessionManager;
use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult};
use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, Router, Scheduler, SchedulerDeps};
@@ -84,6 +85,15 @@ fn resolve_owner_scope_notification_user(
trimmed_option(explicit_user).or_else(|| trimmed_option(owner_fallback))
}
fn is_single_message_repl(message: &IncomingMessage) -> bool {
message.channel == "repl"
&& message
.metadata
.get("single_message_mode")
.and_then(|value| value.as_bool())
.unwrap_or(false)
}
async fn resolve_channel_notification_user(
extension_manager: Option<&Arc<ExtensionManager>>,
channel: Option<&str>,
@@ -172,6 +182,8 @@ pub struct AgentDeps {
/// Resolved LLM backend identifier (e.g., "nearai", "openai", "groq").
/// Used by `/model` persistence to determine which env var to update.
pub llm_backend: String,
/// Per-tenant rate limiting registry (lazily creates rate state per user).
pub tenant_rates: Arc<crate::tenant::TenantRateRegistry>,
}
/// The main agent that coordinates all components.
@@ -234,7 +246,10 @@ impl Agent {
SchedulerDeps {
tools: deps.tools.clone(),
extension_manager: deps.extension_manager.clone(),
store: deps.store.clone(),
store: deps
.store
.as_ref()
.map(|db| crate::tenant::AdminScope::new(Arc::clone(db))),
hooks: deps.hooks.clone(),
},
);
@@ -315,6 +330,50 @@ impl Agent {
&self.deps.cost_guard
}
/// Build a tenant-scoped execution context for the given user.
///
/// This is the standard entry point for per-user operations. The returned
/// [`TenantCtx`] provides a [`TenantScope`] that auto-binds `user_id` on
/// every database operation and a per-user rate limiter.
pub(super) async fn tenant_ctx(&self, user_id: &str) -> crate::tenant::TenantCtx {
let rate = self.deps.tenant_rates.get_or_create(user_id).await;
let store = self
.deps
.store
.as_ref()
.map(|db| crate::tenant::TenantScope::new(user_id, Arc::clone(db)));
// Reuse the owner workspace if user matches, otherwise create per-user.
let workspace = match &self.deps.workspace {
Some(ws) if ws.user_id() == user_id => Some(Arc::clone(ws)),
_ => self
.deps
.store
.as_ref()
.map(|db| Arc::new(Workspace::new_with_db(user_id, Arc::clone(db)))),
};
crate::tenant::TenantCtx::new(
user_id,
store,
workspace,
Arc::clone(&self.deps.cost_guard),
rate,
)
}
/// Get an admin-scoped database accessor for cross-tenant operations.
///
/// Only for system-level components (heartbeat, routine engine, self-repair,
/// scheduler). Handler code should use [`tenant_ctx()`](Self::tenant_ctx) instead.
pub(super) fn admin_store(&self) -> Option<crate::tenant::AdminScope> {
self.deps
.store
.as_ref()
.map(|db| crate::tenant::AdminScope::new(Arc::clone(db)))
}
pub(super) fn skill_registry(&self) -> Option<&Arc<std::sync::RwLock<SkillRegistry>>> {
self.deps.skill_registry.as_ref()
}
@@ -400,8 +459,8 @@ impl Agent {
self.config.stuck_threshold,
self.config.max_repair_attempts,
);
if let Some(ref store) = self.deps.store {
self_repair = self_repair.with_store(Arc::clone(store));
if let Some(admin) = self.admin_store() {
self_repair = self_repair.with_store(admin);
}
if let Some(ref builder) = self.deps.builder {
self_repair = self_repair.with_builder(Arc::clone(builder), Arc::clone(self.tools()));
@@ -508,6 +567,7 @@ impl Agent {
.with_interval(std::time::Duration::from_secs(hb_config.interval_secs));
config.quiet_hours_start = hb_config.quiet_hours_start;
config.quiet_hours_end = hb_config.quiet_hours_end;
config.multi_tenant = hb_config.multi_tenant;
config.timezone = hb_config
.timezone
.clone()
@@ -537,30 +597,52 @@ impl Agent {
.await;
let notify_user = heartbeat_notify_user;
let channels = self.channels.clone();
let is_multi_tenant = hb_config.multi_tenant;
tokio::spawn(async move {
while let Some(response) = notify_rx.recv().await {
// In multi-tenant mode, extract the owning user_id from
// the response metadata so notifications reach the
// correct user rather than the agent's owner.
// This intentionally overrides the configured notify_target
// because each user's heartbeat should notify that user.
let effective_user = if is_multi_tenant {
response
.metadata
.get("owner_id")
.and_then(|v| v.as_str())
.map(String::from)
} else {
None
};
// Try the configured channel first, fall back to
// broadcasting on all channels.
let targeted_ok = if let Some(ref channel) = notify_channel
&& let Some(ref user) = notify_target
{
channels
.broadcast(channel, user, response.clone())
.await
.is_ok()
let targeted_ok = if let Some(ref channel) = notify_channel {
let target = effective_user.as_deref().or(notify_target.as_deref());
if let Some(user) = target {
channels
.broadcast(channel, user, response.clone())
.await
.is_ok()
} else {
false
}
} else {
false
};
if !targeted_ok && let Some(ref user) = notify_user {
let results = channels.broadcast_all(user, response).await;
for (ch, result) in results {
if let Err(e) = result {
tracing::warn!(
"Failed to broadcast heartbeat to {}: {}",
ch,
e
);
if !targeted_ok {
let fallback = effective_user.as_deref().or(notify_user.as_deref());
if let Some(user) = fallback {
let results = channels.broadcast_all(user, response).await;
for (ch, result) in results {
if let Err(e) = result {
tracing::warn!(
"Failed to broadcast heartbeat to {}: {}",
ch,
e
);
}
}
}
}
@@ -573,14 +655,29 @@ impl Agent {
.map(|h| h.to_workspace_config())
.unwrap_or_default();
Some(spawn_heartbeat(
config,
hygiene,
workspace.clone(),
self.cheap_llm().clone(),
Some(notify_tx),
self.store().map(Arc::clone),
))
if config.multi_tenant {
if let Some(admin) = self.admin_store() {
Some(spawn_multi_user_heartbeat(
config,
hygiene,
self.cheap_llm().clone(),
Some(notify_tx),
admin,
))
} else {
tracing::warn!("Multi-tenant heartbeat requires a database store");
None
}
} else {
Some(spawn_heartbeat(
config,
hygiene,
workspace.clone(),
self.cheap_llm().clone(),
Some(notify_tx),
self.admin_store(),
))
}
} else {
tracing::warn!("Heartbeat enabled but no workspace available");
None
@@ -602,7 +699,7 @@ impl Agent {
let engine = Arc::new(RoutineEngine::new(
rt_config.clone(),
Arc::clone(store),
crate::tenant::AdminScope::new(Arc::clone(store)),
self.llm().clone(),
Arc::clone(workspace),
notify_tx,
@@ -1140,9 +1237,14 @@ impl Agent {
&& let Submission::UserInput { ref content } = submission
&& let Some(engine) = self.routine_engine().await
{
let single_message_repl = is_single_message_repl(message);
// Use post-hook content so that BeforeInbound hooks that rewrite
// input are respected by event trigger matching.
let fired = engine.check_event_triggers(message, content).await;
let fired = if single_message_repl {
engine.check_event_triggers_and_wait(message, content).await
} else {
engine.check_event_triggers(message, content).await
};
if fired > 0 {
tracing::debug!(
channel = %message.channel,
@@ -1150,15 +1252,30 @@ impl Agent {
fired,
"Consumed inbound user message with matching event-triggered routine(s)"
);
return Ok(Some(String::new()));
return if single_message_repl {
Ok(None)
} else {
Ok(Some(String::new()))
};
}
}
// Build per-tenant execution context once; threaded through all handlers.
let tenant = self.tenant_ctx(&message.user_id).await;
let session_for_empty_exit = Arc::clone(&session);
// Process based on submission type
let result = match submission {
Submission::UserInput { content } => {
let mut result = self
.process_user_input(message, session.clone(), thread_id, &content)
.process_user_input(
message,
tenant.clone(),
session.clone(),
thread_id,
&content,
)
.await;
// Drain any messages queued during processing.
@@ -1225,7 +1342,13 @@ impl Agent {
let mut queued_msg = message.clone();
queued_msg.attachments.clear();
result = self
.process_user_input(&queued_msg, session.clone(), thread_id, &next_content)
.process_user_input(
&queued_msg,
tenant.clone(),
session.clone(),
thread_id,
&next_content,
)
.await;
// If processing failed, re-queue the drained content so it
@@ -1250,8 +1373,30 @@ impl Agent {
command,
message.channel
);
// /reasoning is special-cased here (not in handle_system_command)
// because it needs the session + thread_id to read turn reasoning
// data, which handle_system_command's signature doesn't provide.
if command == "reasoning" {
let result = self
.handle_reasoning_command(&args, &session, thread_id)
.await;
return match result {
SubmissionResult::Response { content } => Ok(Some(content)),
SubmissionResult::Ok { message } => Ok(message),
SubmissionResult::Error { message } => {
Ok(Some(format!("Error: {}", message)))
}
_ => {
if is_single_message_repl(message) {
Ok(None)
} else {
Ok(Some(String::new()))
}
}
};
}
// Authorization checks (including restart channel check) are enforced in handle_system_command
self.handle_system_command(&command, &args, &message.channel)
self.handle_system_command(&command, &args, &message.channel, &tenant)
.await
}
Submission::Undo => self.process_undo(session, thread_id).await,
@@ -1264,12 +1409,9 @@ impl Agent {
Submission::Summarize => self.process_summarize(session, thread_id).await,
Submission::Suggest => self.process_suggest(session, thread_id).await,
Submission::JobStatus { job_id } => {
self.process_job_status(&message.user_id, job_id.as_deref())
.await
}
Submission::JobCancel { job_id } => {
self.process_job_cancel(&message.user_id, &job_id).await
self.process_job_status(&tenant, job_id.as_deref()).await
}
Submission::JobCancel { job_id } => self.process_job_cancel(&tenant, &job_id).await,
Submission::Quit => return Ok(None),
Submission::SwitchThread { thread_id: target } => {
self.process_switch_thread(message, target).await
@@ -1309,7 +1451,26 @@ impl Agent {
Ok(Some(content))
}
}
SubmissionResult::Ok { message } => Ok(message),
SubmissionResult::Ok {
message: output_message,
} => {
let should_exit =
if output_message.as_deref() == Some("") && is_single_message_repl(message) {
let sess = session_for_empty_exit.lock().await;
sess.threads
.get(&thread_id)
.map(|thread| thread.state != ThreadState::AwaitingApproval)
.unwrap_or(true)
} else {
false
};
if should_exit {
Ok(None)
} else {
Ok(output_message)
}
}
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
SubmissionResult::Interrupted => Ok(Some("Interrupted.".into())),
SubmissionResult::NeedApproval { .. } => {
@@ -1325,7 +1486,7 @@ impl Agent {
#[cfg(test)]
mod tests {
use super::{
chat_tool_execution_metadata, resolve_routine_notification_user,
chat_tool_execution_metadata, is_single_message_repl, resolve_routine_notification_user,
should_fallback_routine_notification, truncate_for_preview,
};
use crate::channels::IncomingMessage;
@@ -1487,4 +1648,17 @@ mod tests {
assert!(should_fallback_routine_notification(&error)); // safety: test-only assertion
}
#[test]
fn single_message_repl_detection_requires_repl_channel_and_metadata_flag() {
let repl = IncomingMessage::new("repl", "owner-scope", "hello")
.with_metadata(serde_json::json!({ "single_message_mode": true }));
let gateway = IncomingMessage::new("gateway", "owner-scope", "hello")
.with_metadata(serde_json::json!({ "single_message_mode": true }));
let plain_repl = IncomingMessage::new("repl", "owner-scope", "hello");
assert!(is_single_message_repl(&repl)); // safety: test-only assertion
assert!(!is_single_message_repl(&gateway)); // safety: test-only assertion
assert!(!is_single_message_repl(&plain_repl)); // safety: test-only assertion
}
}
+126 -1
View File
@@ -10,7 +10,7 @@ use std::borrow::Cow;
use crate::agent::session::PendingApproval;
use crate::error::Error;
use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult};
use crate::llm::{ChatMessage, FinishReason, Reasoning, ReasoningContext, RespondResult};
/// Signal from the delegate indicating how the loop should proceed.
pub enum LoopSignal {
@@ -134,6 +134,9 @@ pub async fn run_agentic_loop(
config: &AgenticLoopConfig,
) -> Result<LoopOutcome, Error> {
let mut consecutive_tool_intent_nudges: u32 = 0;
// Accumulates across all iterations (not reset by text responses) so
// non-consecutive truncations still escalate to force_text.
let mut truncation_count: u32 = 0;
for iteration in 1..=config.max_iterations {
// Check for external signals (stop, cancellation, user messages)
@@ -215,7 +218,35 @@ pub async fn run_agentic_loop(
tool_calls,
content,
} => {
// If the response was truncated, tool call parameters are likely
// incomplete. Discard them and tell the LLM to try a different
// approach rather than executing malformed tool calls.
if output.finish_reason == FinishReason::Length {
truncation_count += 1;
let names: Vec<&str> = tool_calls.iter().map(|tc| tc.name.as_str()).collect();
tracing::warn!(
iteration,
tools = ?names,
truncation_count,
"Discarding truncated tool calls (finish_reason=Length)"
);
if let Some(ref text) = content {
reason_ctx.messages.push(ChatMessage::assistant(text));
}
reason_ctx
.messages
.push(ChatMessage::user(crate::llm::TRUNCATED_TOOL_CALL_NOTICE));
// After repeated truncations, force text-only mode so the LLM
// stops attempting tool calls it can't fit in the output budget.
if truncation_count >= 3 {
reason_ctx.force_text = true;
}
delegate.after_iteration(iteration).await;
continue;
}
consecutive_tool_intent_nudges = 0;
truncation_count = 0;
if let Some(outcome) = delegate
.execute_tool_calls(tool_calls, content, reason_ctx)
@@ -271,6 +302,7 @@ mod tests {
RespondOutput {
result: RespondResult::Text(text.to_string()),
usage: zero_usage(),
finish_reason: FinishReason::Stop,
}
}
@@ -281,6 +313,7 @@ mod tests {
content: None,
},
usage: zero_usage(),
finish_reason: FinishReason::ToolUse,
}
}
@@ -414,6 +447,7 @@ mod tests {
id: "call_1".to_string(),
name: "echo".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
};
let delegate = MockDelegate::new(vec![
tool_calls_output(vec![tool_call]),
@@ -621,4 +655,95 @@ mod tests {
let result = truncate_for_preview("café", 4);
assert_eq!(result, "caf...");
}
#[tokio::test]
async fn test_truncated_tool_calls_discarded_on_length() {
let truncated_tool_call = ToolCall {
id: "call_1".to_string(),
name: "memory_write".to_string(),
arguments: serde_json::json!({}), // empty — truncated
reasoning: None,
};
let truncated_output = RespondOutput {
result: RespondResult::ToolCalls {
tool_calls: vec![truncated_tool_call],
content: Some("I'll write the report.".to_string()),
},
usage: zero_usage(),
finish_reason: FinishReason::Length, // response was truncated
};
let delegate = MockDelegate::new(vec![truncated_output, text_output("Summarized it.")]);
let reasoning = stub_reasoning();
let mut ctx = ReasoningContext::new();
let config = AgenticLoopConfig {
max_iterations: 5,
..Default::default()
};
let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config)
.await
.unwrap();
// Tool calls should NOT have been executed
assert_eq!(delegate.tool_exec_count.load(Ordering::SeqCst), 0);
// The loop should have continued and returned the text response
assert!(matches!(outcome, LoopOutcome::Response(ref t) if t == "Summarized it."));
// A truncation notice should have been injected into context
assert!(
ctx.messages
.iter()
.any(|m| m.role == crate::llm::Role::User && m.content.contains("truncated")),
"Should inject truncation notice into context"
);
// The partial assistant content should have been preserved
assert!(
ctx.messages
.iter()
.any(|m| m.role == crate::llm::Role::Assistant
&& m.content.contains("write the report")),
"Should preserve partial assistant content"
);
}
#[tokio::test]
async fn test_repeated_truncations_force_text_mode() {
let make_truncated = || RespondOutput {
result: RespondResult::ToolCalls {
tool_calls: vec![ToolCall {
id: "call_1".to_string(),
name: "memory_write".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
}],
content: None,
},
usage: zero_usage(),
finish_reason: FinishReason::Length,
};
// Three truncated responses, then a text response
let delegate = MockDelegate::new(vec![
make_truncated(),
make_truncated(),
make_truncated(),
text_output("Gave up on tool calls."),
]);
let reasoning = stub_reasoning();
let mut ctx = ReasoningContext::new();
let config = AgenticLoopConfig {
max_iterations: 5,
..Default::default()
};
let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config)
.await
.unwrap();
assert!(matches!(outcome, LoopOutcome::Response(_)));
assert_eq!(delegate.tool_exec_count.load(Ordering::SeqCst), 0);
// After 3 truncations, force_text should be set
assert!(
ctx.force_text,
"Should escalate to force_text after repeated truncations"
);
}
}
+177 -53
View File
@@ -33,6 +33,7 @@ impl Agent {
&self,
intent: MessageIntent,
message: &IncomingMessage,
tenant: &crate::tenant::TenantCtx,
) -> Result<SubmissionResult, Error> {
// Send thinking status for non-trivial operations
if let MessageIntent::CreateJob { .. } = &intent {
@@ -52,24 +53,18 @@ impl Agent {
description,
category,
} => {
self.handle_create_job(&message.user_id, title, description, category)
self.handle_create_job(tenant, title, description, category)
.await?
}
MessageIntent::CheckJobStatus { job_id } => {
self.handle_check_status(&message.user_id, job_id).await?
}
MessageIntent::CancelJob { job_id } => {
self.handle_cancel_job(&message.user_id, &job_id).await?
}
MessageIntent::ListJobs { filter } => {
self.handle_list_jobs(&message.user_id, filter).await?
}
MessageIntent::HelpJob { job_id } => {
self.handle_help_job(&message.user_id, &job_id).await?
self.handle_check_status(tenant, job_id).await?
}
MessageIntent::CancelJob { job_id } => self.handle_cancel_job(tenant, &job_id).await?,
MessageIntent::ListJobs { filter } => self.handle_list_jobs(tenant, filter).await?,
MessageIntent::HelpJob { job_id } => self.handle_help_job(tenant, &job_id).await?,
MessageIntent::Command { command, args } => {
match self
.handle_command(&command, &args, &message.channel)
.handle_command(&command, &args, &message.channel, tenant)
.await?
{
Some(s) => s,
@@ -83,14 +78,14 @@ impl Agent {
async fn handle_create_job(
&self,
user_id: &str,
tenant: &crate::tenant::TenantCtx,
title: String,
description: String,
category: Option<String>,
) -> Result<String, Error> {
let job_id = self
.scheduler
.dispatch_job(user_id, &title, &description, None)
.dispatch_job(tenant.user_id(), &title, &description, None)
.await?;
// Set the dedicated category field (not stored in metadata)
@@ -113,7 +108,7 @@ impl Agent {
async fn handle_check_status(
&self,
user_id: &str,
tenant: &crate::tenant::TenantCtx,
job_id: Option<String>,
) -> Result<String, Error> {
match job_id {
@@ -122,7 +117,8 @@ impl Agent {
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
// Try DB first for persistent state, fall back to ContextManager.
if let Some(store) = self.store()
// TenantScope.get_job() auto-filters by ownership — no manual check needed.
if let Some(store) = tenant.store()
&& let Ok(Some(ctx)) = store.get_job(uuid).await
{
return Ok(format!(
@@ -138,7 +134,7 @@ impl Agent {
}
let ctx = self.context_manager.get_context(uuid).await?;
if ctx.user_id != user_id {
if ctx.user_id != tenant.user_id() {
return Err(crate::error::JobError::NotFound { id: uuid }.into());
}
@@ -155,7 +151,8 @@ impl Agent {
}
None => {
// Show summary from DB for consistency with Jobs tab.
if let Some(store) = self.store() {
// TenantScope methods auto-scope to user — no user_id parameter needed.
if let Some(store) = tenant.store() {
let mut total = 0;
let mut in_progress = 0;
let mut completed = 0;
@@ -183,7 +180,7 @@ impl Agent {
}
// Fallback to ContextManager if no DB.
let summary = self.context_manager.summary_for(user_id).await;
let summary = self.context_manager.summary_for(tenant.user_id()).await;
Ok(format!(
"Jobs summary: Total: {} In Progress: {} Completed: {} Failed: {} Stuck: {}",
summary.total,
@@ -196,19 +193,24 @@ impl Agent {
}
}
async fn handle_cancel_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
async fn handle_cancel_job(
&self,
tenant: &crate::tenant::TenantCtx,
job_id: &str,
) -> Result<String, Error> {
let uuid = Uuid::parse_str(job_id)
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
let ctx = self.context_manager.get_context(uuid).await?;
if ctx.user_id != user_id {
if ctx.user_id != tenant.user_id() {
return Err(crate::error::JobError::NotFound { id: uuid }.into());
}
self.scheduler.stop(uuid).await?;
// Also update DB so the Jobs tab reflects cancellation immediately.
if let Some(store) = self.store()
// Use TenantScope — ownership already verified above.
if let Some(store) = tenant.store()
&& let Err(e) = store
.update_job_status(uuid, JobState::Cancelled, Some("Cancelled by user"))
.await
@@ -221,11 +223,12 @@ impl Agent {
async fn handle_list_jobs(
&self,
user_id: &str,
tenant: &crate::tenant::TenantCtx,
_filter: Option<String>,
) -> Result<String, Error> {
// List from DB for consistency with Jobs tab.
if let Some(store) = self.store() {
// TenantScope methods auto-scope to user.
if let Some(store) = tenant.store() {
let agent_jobs = match store.list_agent_jobs().await {
Ok(jobs) => jobs,
Err(e) => {
@@ -256,7 +259,7 @@ impl Agent {
}
// Fallback to ContextManager if no DB.
let jobs = self.context_manager.all_jobs_for(user_id).await;
let jobs = self.context_manager.all_jobs_for(tenant.user_id()).await;
if jobs.is_empty() {
return Ok("No jobs found.".to_string());
}
@@ -270,12 +273,16 @@ impl Agent {
Ok(output)
}
async fn handle_help_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
async fn handle_help_job(
&self,
tenant: &crate::tenant::TenantCtx,
job_id: &str,
) -> Result<String, Error> {
let uuid = Uuid::parse_str(job_id)
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
let ctx = self.context_manager.get_context(uuid).await?;
if ctx.user_id != user_id {
if ctx.user_id != tenant.user_id() {
return Err(crate::error::JobError::NotFound { id: uuid }.into());
}
@@ -308,11 +315,11 @@ impl Agent {
/// Show job status inline — either all jobs (no id) or a specific job.
pub(super) async fn process_job_status(
&self,
user_id: &str,
tenant: &crate::tenant::TenantCtx,
job_id: Option<&str>,
) -> Result<SubmissionResult, Error> {
match self
.handle_check_status(user_id, job_id.map(|s| s.to_string()))
.handle_check_status(tenant, job_id.map(|s| s.to_string()))
.await
{
Ok(text) => Ok(SubmissionResult::response(text)),
@@ -323,10 +330,10 @@ impl Agent {
/// Cancel a job by ID.
pub(super) async fn process_job_cancel(
&self,
user_id: &str,
tenant: &crate::tenant::TenantCtx,
job_id: &str,
) -> Result<SubmissionResult, Error> {
match self.handle_cancel_job(user_id, job_id).await {
match self.handle_cancel_job(tenant, job_id).await {
Ok(text) => Ok(SubmissionResult::response(text)),
Err(e) => Ok(SubmissionResult::error(format!("Cancel error: {}", e))),
}
@@ -465,12 +472,101 @@ impl Agent {
}
}
/// Handle `/reasoning [N|all]` — show reasoning history for the active thread.
pub(super) async fn handle_reasoning_command(
&self,
args: &[String],
session: &Arc<Mutex<Session>>,
thread_id: Uuid,
) -> SubmissionResult {
// Clone the turn data we need, then drop the session lock.
let turns_snapshot: Vec<(
usize,
Option<String>,
Vec<crate::agent::session::TurnToolCall>,
)>;
{
let sess = session.lock().await;
let thread = match sess.threads.get(&thread_id) {
Some(t) => t,
None => return SubmissionResult::error("No active thread."),
};
if thread.turns.is_empty() {
return SubmissionResult::ok_with_message("No turns yet.");
}
// Parse argument: default=last turn, "all"=all turns, N=specific turn (1-based).
let selected: Vec<&crate::agent::session::Turn> = match args.first().map(|s| s.as_str())
{
Some("all") => thread.turns.iter().collect(),
Some(n) => match n.parse::<usize>() {
Ok(0) => return SubmissionResult::error("Turn numbers start at 1."),
Ok(num) if num > thread.turns.len() => {
return SubmissionResult::error(format!(
"Turn {} does not exist (max: {}).",
num,
thread.turns.len()
));
}
Ok(num) => vec![&thread.turns[num - 1]],
Err(_) => return SubmissionResult::error("Usage: /reasoning [N|all]"),
},
None => {
// Default: last turn that has tool calls
match thread.turns.iter().rev().find(|t| !t.tool_calls.is_empty()) {
Some(t) => vec![t],
None => {
return SubmissionResult::ok_with_message("No turns with tool calls.");
}
}
}
};
turns_snapshot = selected
.into_iter()
.map(|t| (t.turn_number, t.narrative.clone(), t.tool_calls.clone()))
.collect();
}
// Session lock is now dropped — format output without holding it.
let mut output = String::new();
for (turn_number, narrative, tool_calls) in &turns_snapshot {
output.push_str(&format!("--- Turn {} ---\n", turn_number + 1));
if let Some(narrative) = narrative {
output.push_str(&format!("Reasoning: {}\n", narrative));
}
if tool_calls.is_empty() {
output.push_str(" (no tool calls)\n");
} else {
for tc in tool_calls {
let status = if tc.error.is_some() {
"error"
} else if tc.result.is_some() {
"ok"
} else {
"pending"
};
output.push_str(&format!(" {} [{}]", tc.name, status));
if let Some(ref rationale) = tc.rationale {
output.push_str(&format!("{}", rationale));
}
output.push('\n');
}
}
output.push('\n');
}
SubmissionResult::response(output.trim_end())
}
/// Handle system commands that bypass thread-state checks entirely.
pub(super) async fn handle_system_command(
&self,
command: &str,
args: &[String],
channel: &str,
tenant: &crate::tenant::TenantCtx,
) -> Result<SubmissionResult, Error> {
match command {
"help" => Ok(SubmissionResult::response(concat!(
@@ -480,6 +576,7 @@ impl Agent {
" /version Show version info\n",
" /tools List available tools\n",
" /debug Toggle debug mode\n",
" /reasoning [N|all] Show agent reasoning for turns\n",
" /ping Connectivity check\n",
"\n",
"Jobs:\n",
@@ -663,19 +760,32 @@ impl Agent {
}
}
match self.llm().set_model(requested) {
Ok(()) => {
// Persist the model choice so it survives restarts.
self.persist_selected_model(requested).await;
Ok(SubmissionResult::response(format!(
"Switched model to: {}",
requested
)))
if self.config.multi_tenant {
// Multi-tenant: only persist to per-user DB settings.
// Do NOT call set_model() on the shared provider — that
// would change the default for all users. The per-request
// model_override in the dispatcher reads from the same
// "selected_model" setting and applies it per-user.
self.persist_selected_model(tenant, requested).await;
Ok(SubmissionResult::response(format!(
"Model preference set to: {} (per-user)",
requested
)))
} else {
match self.llm().set_model(requested) {
Ok(()) => {
// Persist the model choice so it survives restarts.
self.persist_selected_model(tenant, requested).await;
Ok(SubmissionResult::response(format!(
"Switched model to: {}",
requested
)))
}
Err(e) => Ok(SubmissionResult::error(format!(
"Failed to switch model: {}",
e
))),
}
Err(e) => Ok(SubmissionResult::error(format!(
"Failed to switch model: {}",
e
))),
}
}
}
@@ -817,10 +927,14 @@ impl Agent {
command: &str,
args: &[String],
channel: &str,
tenant: &crate::tenant::TenantCtx,
) -> Result<Option<String>, Error> {
// System commands are now handled directly via Submission::SystemCommand,
// but the router may still send us unknown /commands.
match self.handle_system_command(command, args, channel).await? {
match self
.handle_system_command(command, args, channel, tenant)
.await?
{
SubmissionResult::Response { content } => Ok(Some(content)),
SubmissionResult::Ok { message } => Ok(message),
SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))),
@@ -832,23 +946,33 @@ impl Agent {
///
/// Best-effort: logs warnings on failure but does not propagate errors,
/// since the in-memory model switch already succeeded.
async fn persist_selected_model(&self, model: &str) {
// 1. Persist to DB if available.
if let Some(store) = self.store() {
///
/// In multi-tenant mode, only the per-user DB setting is written — global
/// .env and TOML files are shared across users and must not be mutated.
async fn persist_selected_model(&self, tenant: &crate::tenant::TenantCtx, model: &str) {
// 1. Persist to DB if available (per-user scoped via TenantScope).
if let Some(store) = tenant.store() {
let value = serde_json::Value::String(model.to_string());
if let Err(e) = store
.set_setting(self.owner_id(), "selected_model", &value)
.await
{
if let Err(e) = store.set_setting("selected_model", &value).await {
tracing::warn!("Failed to persist model to DB: {}", e);
} else {
tracing::debug!("Persisted selected_model to DB: {}", model);
tracing::debug!(
user_id = tenant.user_id(),
"Persisted selected_model to DB: {}",
model
);
}
} else {
tracing::warn!("No database store available — model choice will not persist to DB");
}
// 2. Update .env and TOML config file (sync I/O in spawn_blocking).
// 2. In multi-tenant mode, skip .env/TOML writes — these are global
// files shared by all users. The per-user DB setting is sufficient.
if self.config.multi_tenant {
return;
}
// 3. Update .env and TOML config file (sync I/O in spawn_blocking).
let model_owned = model.to_string();
let backend = self.deps.llm_backend.clone();
if let Err(e) = tokio::task::spawn_blocking(move || {
+236 -3
View File
@@ -21,6 +21,9 @@ pub struct CostGuardConfig {
pub max_cost_per_day_cents: Option<u64>,
/// Maximum LLM calls per hour. None = unlimited.
pub max_actions_per_hour: Option<u64>,
/// Maximum spend per user per day in cents. None = unlimited.
/// Applied independently per user alongside the global budget.
pub max_cost_per_user_per_day_cents: Option<u64>,
}
/// Error returned when a cost limit is exceeded.
@@ -30,6 +33,12 @@ pub enum CostLimitExceeded {
DailyBudget { spent_cents: u64, limit_cents: u64 },
/// Hourly action rate limit reached.
HourlyRate { actions: u64, limit: u64 },
/// Per-user daily spending cap reached.
UserDailyBudget {
user_id: String,
spent_cents: u64,
limit_cents: u64,
},
}
impl std::fmt::Display for CostLimitExceeded {
@@ -49,6 +58,17 @@ impl std::fmt::Display for CostLimitExceeded {
"Hourly action limit exceeded: {} actions of {} allowed per hour",
actions, limit
),
Self::UserDailyBudget {
user_id,
spent_cents,
limit_cents,
} => write!(
f,
"User '{}' daily cost limit exceeded: spent ${:.2} of ${:.2} allowed",
user_id,
*spent_cents as f64 / 100.0,
*limit_cents as f64 / 100.0
),
}
}
}
@@ -78,6 +98,9 @@ pub struct CostGuard {
/// Per-model token usage since startup.
model_tokens: Mutex<HashMap<String, ModelTokens>>,
/// Per-user daily cost tracking. Each entry resets independently at midnight UTC.
per_user_daily_cost: Mutex<HashMap<String, DailyCost>>,
}
struct DailyCost {
@@ -97,6 +120,7 @@ impl CostGuard {
action_window: Mutex::new(VecDeque::new()),
budget_exceeded: AtomicBool::new(false),
model_tokens: Mutex::new(HashMap::new()),
per_user_daily_cost: Mutex::new(HashMap::new()),
}
}
@@ -203,6 +227,11 @@ impl CostGuard {
daily.reset_date = today;
self.budget_exceeded.store(false, Ordering::Relaxed);
tracing::info!("Cost guard: daily counter reset for {}", today);
// Prune per-user entries from previous days to prevent
// unbounded HashMap growth in long-lived deployments.
let mut per_user = self.per_user_daily_cost.lock().await;
per_user.retain(|_, entry| entry.reset_date == today);
}
daily.total += cost;
@@ -248,6 +277,85 @@ impl CostGuard {
cost
}
/// Record an LLM call with per-user attribution.
///
/// Delegates to `record_llm_call` for global tracking, then additionally
/// records the cost against the user's daily budget.
#[allow(clippy::too_many_arguments)]
pub async fn record_llm_call_for_user(
&self,
user_id: &str,
model: &str,
input_tokens: u32,
output_tokens: u32,
cache_read_input_tokens: u32,
cache_creation_input_tokens: u32,
cache_read_discount: Decimal,
cache_write_multiplier: Decimal,
cost_per_token: Option<(Decimal, Decimal)>,
) -> Decimal {
let cost = self
.record_llm_call(
model,
input_tokens,
output_tokens,
cache_read_input_tokens,
cache_creation_input_tokens,
cache_read_discount,
cache_write_multiplier,
cost_per_token,
)
.await;
// Track per-user daily cost
{
let today = chrono::Utc::now().date_naive();
let mut per_user = self.per_user_daily_cost.lock().await;
let entry = per_user
.entry(user_id.to_string())
.or_insert_with(|| DailyCost {
total: Decimal::ZERO,
reset_date: today,
});
if today != entry.reset_date {
entry.total = Decimal::ZERO;
entry.reset_date = today;
}
entry.total += cost;
}
cost
}
/// Check whether the next action is allowed for a specific user.
///
/// Checks the global limits first (via `check_allowed`), then additionally
/// checks the per-user daily budget if configured.
pub async fn check_allowed_for_user(&self, user_id: &str) -> Result<(), CostLimitExceeded> {
// Check global limits first
self.check_allowed().await?;
// Check per-user daily budget
if let Some(limit_cents) = self.config.max_cost_per_user_per_day_cents {
let today = chrono::Utc::now().date_naive();
let per_user = self.per_user_daily_cost.lock().await;
if let Some(entry) = per_user.get(user_id)
&& entry.reset_date == today
{
let spent_cents = to_cents(entry.total);
if spent_cents >= limit_cents {
return Err(CostLimitExceeded::UserDailyBudget {
user_id: user_id.to_string(),
spent_cents,
limit_cents,
});
}
}
}
Ok(())
}
/// Current daily spend in USD (as Decimal).
pub async fn daily_spend(&self) -> Decimal {
let daily = self.daily_cost.lock().await;
@@ -259,6 +367,16 @@ impl CostGuard {
}
}
/// Current daily spend for a specific user in USD (as Decimal).
pub async fn daily_spend_for_user(&self, user_id: &str) -> Decimal {
let today = chrono::Utc::now().date_naive();
let per_user = self.per_user_daily_cost.lock().await;
match per_user.get(user_id) {
Some(entry) if entry.reset_date == today => entry.total,
_ => Decimal::ZERO,
}
}
/// Number of actions in the current hourly window.
pub async fn actions_this_hour(&self) -> u64 {
let mut window = self.action_window.lock().await;
@@ -314,7 +432,7 @@ mod tests {
async fn test_daily_budget_enforcement() {
let guard = CostGuard::new(CostGuardConfig {
max_cost_per_day_cents: Some(1), // $0.01 limit
max_actions_per_hour: None,
..CostGuardConfig::default()
});
// First call allowed
@@ -350,8 +468,8 @@ mod tests {
#[tokio::test]
async fn test_hourly_rate_enforcement() {
let guard = CostGuard::new(CostGuardConfig {
max_cost_per_day_cents: None,
max_actions_per_hour: Some(3),
..CostGuardConfig::default()
});
// First 3 actions allowed
@@ -633,8 +751,8 @@ mod tests {
// A fresh CostGuard with rate limits should not panic even if
// checked_sub returns None (simulating short uptime).
let guard = CostGuard::new(CostGuardConfig {
max_cost_per_day_cents: None,
max_actions_per_hour: Some(100),
..CostGuardConfig::default()
});
// These must not panic regardless of system uptime
@@ -656,4 +774,119 @@ mod tests {
let result = Instant::now().checked_sub(std::time::Duration::MAX);
assert!(result.is_none());
}
#[tokio::test]
async fn test_per_user_daily_budget_enforcement() {
let guard = CostGuard::new(CostGuardConfig {
max_cost_per_day_cents: None,
max_actions_per_hour: None,
max_cost_per_user_per_day_cents: Some(1), // $0.01 per user
});
// Both users initially allowed
assert!(guard.check_allowed_for_user("alice").await.is_ok());
assert!(guard.check_allowed_for_user("bob").await.is_ok());
// Alice makes an expensive call
guard
.record_llm_call_for_user(
"alice",
"gpt-4o",
10_000,
10_000,
0,
0,
Decimal::ONE,
Decimal::ONE,
None,
)
.await;
// Alice should be blocked, Bob should still be allowed
let result = guard.check_allowed_for_user("alice").await;
assert!(result.is_err());
match result.unwrap_err() {
CostLimitExceeded::UserDailyBudget {
user_id,
limit_cents,
..
} => {
assert_eq!(user_id, "alice");
assert_eq!(limit_cents, 1);
}
other => panic!("Expected UserDailyBudget, got {:?}", other),
}
assert!(guard.check_allowed_for_user("bob").await.is_ok());
}
#[tokio::test]
async fn test_per_user_daily_spend_tracking() {
let guard = CostGuard::new(CostGuardConfig::default());
assert_eq!(guard.daily_spend_for_user("alice").await, Decimal::ZERO);
assert_eq!(guard.daily_spend_for_user("bob").await, Decimal::ZERO);
let cost = guard
.record_llm_call_for_user(
"alice",
"gpt-4o",
1000,
500,
0,
0,
Decimal::ONE,
Decimal::ONE,
None,
)
.await;
assert_eq!(guard.daily_spend_for_user("alice").await, cost);
assert_eq!(guard.daily_spend_for_user("bob").await, Decimal::ZERO);
// Global spend should also be tracked
assert_eq!(guard.daily_spend().await, cost);
}
#[tokio::test]
async fn test_per_user_budget_independent_of_global() {
let guard = CostGuard::new(CostGuardConfig {
max_cost_per_day_cents: Some(100_000), // $1000 global limit
max_actions_per_hour: None,
max_cost_per_user_per_day_cents: Some(1), // $0.01 per user
});
// User hits their personal limit
guard
.record_llm_call_for_user(
"alice",
"gpt-4o",
10_000,
10_000,
0,
0,
Decimal::ONE,
Decimal::ONE,
None,
)
.await;
// Alice blocked by per-user limit, not global
assert!(guard.check_allowed_for_user("alice").await.is_err());
// Global limit is far from reached
assert!(guard.check_allowed().await.is_ok());
// Bob is unaffected
assert!(guard.check_allowed_for_user("bob").await.is_ok());
}
#[test]
fn test_user_cost_limit_display() {
let limit = CostLimitExceeded::UserDailyBudget {
user_id: "alice".to_string(),
spent_cents: 150,
limit_cents: 100,
};
let msg = limit.to_string();
assert!(msg.contains("alice"));
assert!(msg.contains("$1.50"));
assert!(msg.contains("$1.00"));
}
}
+140 -15
View File
@@ -42,6 +42,7 @@ impl Agent {
pub(super) async fn run_agentic_loop(
&self,
message: &IncomingMessage,
tenant: crate::tenant::TenantCtx,
session: Arc<Mutex<Session>>,
thread_id: Uuid,
initial_messages: Vec<ChatMessage>,
@@ -63,7 +64,12 @@ impl Agent {
);
let system_prompt = if let Some(ws) = self.workspace() {
match ws
let scoped_workspace = if ws.user_id() == message.user_id {
Arc::clone(ws)
} else {
Arc::new(ws.scoped_to_user(&message.user_id))
};
match scoped_workspace
.system_prompt_for_context_tz(is_group_chat, user_tz)
.await
{
@@ -163,6 +169,7 @@ impl Agent {
let delegate = ChatDelegate {
agent: self,
tenant,
session: session.clone(),
thread_id,
message,
@@ -235,6 +242,7 @@ impl Agent {
/// auth intercept, and cost tracking.
struct ChatDelegate<'a> {
agent: &'a Agent,
tenant: crate::tenant::TenantCtx,
session: Arc<Mutex<Session>>,
thread_id: Uuid,
message: &'a IncomingMessage,
@@ -298,6 +306,8 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
// Update context for this iteration
reason_ctx.available_tools = tool_defs;
// Preserve force_text if already set (e.g. by truncation escalation).
let force_text = force_text || reason_ctx.force_text;
reason_ctx.system_prompt = Some(if force_text {
self.cached_prompt_no_tools.clone()
} else {
@@ -331,8 +341,8 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
reason_ctx: &mut ReasoningContext,
iteration: usize,
) -> Result<crate::llm::RespondOutput, Error> {
// Enforce cost guardrails before the LLM call
if let Err(limit) = self.agent.cost_guard().check_allowed().await {
// Enforce cost guardrails before the LLM call (global + per-user)
if let Err(limit) = self.tenant.check_cost_allowed().await {
return Err(crate::error::LlmError::InvalidResponse {
provider: "agent".to_string(),
reason: limit.to_string(),
@@ -340,6 +350,21 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
.into());
}
// Apply per-user model override from settings (first iteration only
// to avoid repeated DB lookups within the same agentic loop).
// Uses "selected_model" — the same key the /model command persists to
// via SettingsStore (per-user scoped via TenantScope).
if iteration == 0
&& let Some(store) = self.tenant.store()
&& let Ok(Some(value)) = store.get_setting("selected_model").await
&& let Some(model) = value.as_str()
{
let model = model.trim();
if !model.is_empty() {
reason_ctx.model_override = Some(model.to_string());
}
}
let output = match reasoning.respond_with_tools(reason_ctx).await {
Ok(output) => output,
Err(crate::error::LlmError::ContextLengthExceeded { used, limit }) => {
@@ -374,13 +399,22 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
Err(e) => return Err(e.into()),
};
// Record cost and track token usage
let model_name = self.agent.llm().active_model_name();
// Record cost and track token usage (global + per-user).
// When a model override is active, use the override name for attribution
// and let CostGuard look up pricing via costs::model_cost() instead of
// using the default provider's cost_per_token (which reflects the wrong model).
let (model_name, cost_per_token) = if let Some(ref ovr) = reason_ctx.model_override {
(ovr.clone(), None)
} else {
(
self.agent.llm().active_model_name(),
Some(self.agent.llm().cost_per_token()),
)
};
let read_discount = self.agent.llm().cache_read_discount();
let write_multiplier = self.agent.llm().cache_write_multiplier();
let call_cost = self
.agent
.cost_guard()
.tenant
.record_llm_call(
&model_name,
output.usage.input_tokens,
@@ -389,7 +423,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
output.usage.cache_creation_input_tokens,
read_discount,
write_multiplier,
Some(self.agent.llm().cost_per_token()),
cost_per_token,
)
.await;
tracing::debug!(
@@ -420,6 +454,19 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
content: Option<String>,
reason_ctx: &mut ReasoningContext,
) -> Result<Option<LoopOutcome>, Error> {
// Extract and sanitize the narrative before consuming `content`.
let narrative = content
.as_deref()
.filter(|c| !c.trim().is_empty())
.map(|c| {
let sanitized = self
.agent
.safety()
.sanitize_tool_output("agent_narrative", c);
sanitized.content
})
.filter(|c| !c.trim().is_empty());
// Add the assistant message with tool_calls to context.
// OpenAI protocol requires this before tool-result messages.
reason_ctx
@@ -440,6 +487,41 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
)
.await;
// Build per-tool decisions for the reasoning update.
// Sanitize each rationale through SafetyLayer (parity with JobDelegate).
let decisions: Vec<crate::channels::ToolDecision> = tool_calls
.iter()
.filter_map(|tc| {
tc.reasoning.as_ref().map(|r| {
let sanitized = self
.agent
.safety()
.sanitize_tool_output("tool_rationale", r)
.content;
crate::channels::ToolDecision {
tool_name: tc.name.clone(),
rationale: sanitized,
}
})
})
.collect();
// Emit reasoning update to channels.
if narrative.is_some() || !decisions.is_empty() {
let _ = self
.agent
.channels
.send_status(
&self.message.channel,
StatusUpdate::ReasoningUpdate {
narrative: narrative.clone().unwrap_or_default(),
decisions: decisions.clone(),
},
&self.message.metadata,
)
.await;
}
// Record tool calls in the thread with sensitive params redacted.
{
let mut redacted_args: Vec<serde_json::Value> = Vec::with_capacity(tool_calls.len());
@@ -455,8 +537,23 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
if let Some(thread) = sess.threads.get_mut(&self.thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
// Set turn-level narrative.
if turn.narrative.is_none() {
turn.narrative = narrative;
}
for (tc, safe_args) in tool_calls.iter().zip(redacted_args) {
turn.record_tool_call(&tc.name, safe_args);
let sanitized_rationale = tc.reasoning.as_ref().map(|r| {
self.agent
.safety()
.sanitize_tool_output("tool_rationale", r)
.content
});
turn.record_tool_call_with_reasoning(
&tc.name,
safe_args,
sanitized_rationale,
Some(tc.id.clone()),
);
}
}
}
@@ -726,7 +823,7 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
if let Some(thread) = sess.threads.get_mut(&self.thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
turn.record_tool_error(error_msg.clone());
turn.record_tool_error_for(&tc.id, error_msg.clone());
}
}
reason_ctx
@@ -852,16 +949,19 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
Err(e) => format!("Tool '{}' failed: {}", tc.name, e),
};
// Record sanitized result in thread
// Record sanitized result in thread (identity-based matching).
{
let mut sess = self.session.lock().await;
if let Some(thread) = sess.threads.get_mut(&self.thread_id)
&& let Some(turn) = thread.last_turn_mut()
{
if is_tool_error {
turn.record_tool_error(result_content.clone());
turn.record_tool_error_for(&tc.id, result_content.clone());
} else {
turn.record_tool_result(serde_json::json!(result_content));
turn.record_tool_result_for(
&tc.id,
serde_json::json!(result_content),
);
}
}
}
@@ -1234,6 +1334,7 @@ mod tests {
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
builder: None,
llm_backend: "nearai".to_string(),
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
};
Agent::new(
@@ -1249,10 +1350,14 @@ mod tests {
allow_local_tools: false,
max_cost_per_day_cents: None,
max_actions_per_hour: None,
max_cost_per_user_per_day_cents: None,
max_tool_iterations: 50,
auto_approve_tools: false,
default_timezone: "UTC".to_string(),
max_tokens_per_job: 0,
multi_tenant: false,
max_llm_concurrent_per_user: None,
max_jobs_concurrent_per_user: None,
},
deps,
Arc::new(ChannelManager::new()),
@@ -1462,11 +1567,13 @@ mod tests {
id: "call_2".to_string(),
name: "http".to_string(),
arguments: serde_json::json!({"url": "https://example.com"}),
reasoning: None,
},
ToolCall {
id: "call_3".to_string(),
name: "echo".to_string(),
arguments: serde_json::json!({"message": "done"}),
reasoning: None,
},
],
user_timezone: None,
@@ -1652,6 +1759,7 @@ mod tests {
id: "call_1".to_string(),
name: "echo".to_string(),
arguments: serde_json::json!({"message": "hi"}),
reasoning: None,
}],
),
ChatMessage::tool_result("call_1", "echo", "hi"),
@@ -1744,11 +1852,13 @@ mod tests {
id: "c1".to_string(),
name: "http".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
},
ToolCall {
id: "c2".to_string(),
name: "echo".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
},
],
),
@@ -1782,6 +1892,7 @@ mod tests {
id: "c1".to_string(),
name: "echo".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
}],
),
ChatMessage::tool_result("c1", "echo", "done"),
@@ -1912,6 +2023,7 @@ mod tests {
id: crate::llm::generate_tool_call_id(0, 0),
name: "echo".to_string(),
arguments: serde_json::json!({"message": "looping"}),
reasoning: None,
}],
input_tokens: 0,
output_tokens: 5,
@@ -2065,6 +2177,7 @@ mod tests {
id: crate::llm::generate_tool_call_id(0, 0),
name: "nonexistent_tool".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
}],
input_tokens: 0,
output_tokens: 5,
@@ -2102,6 +2215,7 @@ mod tests {
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
builder: None,
llm_backend: "nearai".to_string(),
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
};
Agent::new(
@@ -2117,10 +2231,14 @@ mod tests {
allow_local_tools: false,
max_cost_per_day_cents: None,
max_actions_per_hour: None,
max_cost_per_user_per_day_cents: None,
max_tool_iterations,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
max_tokens_per_job: 0,
multi_tenant: false,
max_llm_concurrent_per_user: None,
max_jobs_concurrent_per_user: None,
},
deps,
Arc::new(ChannelManager::new()),
@@ -2155,13 +2273,14 @@ mod tests {
let message = IncomingMessage::new("test", "test-user", "do something");
let initial_messages = vec![ChatMessage::user("do something")];
let tenant = agent.tenant_ctx("test-user").await;
// The dispatcher must terminate within 5 seconds. If there is an
// infinite loop bug (e.g., index not advancing on tool failure), the
// timeout will fire and the test will fail.
let result = tokio::time::timeout(
Duration::from_secs(5),
agent.run_agentic_loop(&message, session, thread_id, initial_messages),
agent.run_agentic_loop(&message, tenant, session, thread_id, initial_messages),
)
.await;
@@ -2223,6 +2342,7 @@ mod tests {
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
builder: None,
llm_backend: "nearai".to_string(),
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
};
Agent::new(
@@ -2238,10 +2358,14 @@ mod tests {
allow_local_tools: false,
max_cost_per_day_cents: None,
max_actions_per_hour: None,
max_cost_per_user_per_day_cents: None,
max_tool_iterations: max_iter,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
max_tokens_per_job: 0,
multi_tenant: false,
max_llm_concurrent_per_user: None,
max_jobs_concurrent_per_user: None,
},
deps,
Arc::new(ChannelManager::new()),
@@ -2261,13 +2385,14 @@ mod tests {
let message = IncomingMessage::new("test", "test-user", "keep calling tools");
let initial_messages = vec![ChatMessage::user("keep calling tools")];
let tenant = agent.tenant_ctx("test-user").await;
// Even with an LLM that always wants to call tools, the dispatcher
// must terminate within the timeout thanks to force_text at
// max_tool_iterations.
let result = tokio::time::timeout(
Duration::from_secs(5),
agent.run_agentic_loop(&message, session, thread_id, initial_messages),
agent.run_agentic_loop(&message, tenant, session, thread_id, initial_messages),
)
.await;
+184 -7
View File
@@ -31,8 +31,8 @@ use chrono_tz::Tz;
use tokio::sync::mpsc;
use crate::channels::OutgoingResponse;
use crate::db::Database;
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
use crate::tenant::AdminScope;
use crate::workspace::Workspace;
use crate::workspace::hygiene::HygieneConfig;
@@ -57,6 +57,9 @@ pub struct HeartbeatConfig {
pub quiet_hours_end: Option<u32>,
/// Timezone for fire_at and quiet hours evaluation (IANA name).
pub timezone: Option<String>,
/// When true, cycle through all users with routines instead of
/// running heartbeat for a single user. Requires a database store.
pub multi_tenant: bool,
}
impl Default for HeartbeatConfig {
@@ -71,6 +74,7 @@ impl Default for HeartbeatConfig {
quiet_hours_start: None,
quiet_hours_end: None,
timezone: None,
multi_tenant: false,
}
}
}
@@ -178,7 +182,7 @@ pub struct HeartbeatRunner {
workspace: Arc<Workspace>,
llm: Arc<dyn LlmProvider>,
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
store: Option<Arc<dyn Database>>,
store: Option<AdminScope>,
consecutive_failures: u32,
}
@@ -207,8 +211,8 @@ impl HeartbeatRunner {
self
}
/// Set the database store for persistent heartbeat conversations.
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
/// Set the admin-scoped database store for persistent heartbeat conversations.
pub fn with_store(mut self, store: AdminScope) -> Self {
self.store = Some(store);
self
}
@@ -396,7 +400,7 @@ impl HeartbeatRunner {
}
/// Send a notification about heartbeat findings.
async fn send_notification(&self, message: &str) {
pub(crate) async fn send_notification(&self, message: &str) {
let Some(ref tx) = self.response_tx else {
tracing::debug!("No response channel configured for heartbeat notifications");
return;
@@ -493,7 +497,7 @@ pub fn spawn_heartbeat(
workspace: Arc<Workspace>,
llm: Arc<dyn LlmProvider>,
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
store: Option<Arc<dyn Database>>,
store: Option<AdminScope>,
) -> tokio::task::JoinHandle<()> {
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm);
if let Some(tx) = response_tx {
@@ -508,6 +512,179 @@ pub fn spawn_heartbeat(
})
}
/// Spawn a multi-user heartbeat runner that cycles through all users that
/// own routines (enabled or not). Each tick, it queries the DB for distinct
/// user_ids, creates a per-user workspace, and runs a heartbeat check for
/// each user concurrently. Per-user failure counts are tracked independently.
pub fn spawn_multi_user_heartbeat(
config: HeartbeatConfig,
hygiene_config: HygieneConfig,
llm: Arc<dyn LlmProvider>,
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
store: AdminScope,
) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
if !config.enabled {
tracing::info!("Multi-user heartbeat is disabled");
return;
}
let mut tick_interval = if config.fire_at.is_none() {
let mut iv = tokio::time::interval(config.interval);
iv.tick().await; // skip immediate tick
Some(iv)
} else {
None
};
// Track consecutive failures per user so we can disable heartbeat
// for persistently-failing users (same semantics as single-user mode).
let mut user_failures: std::collections::HashMap<String, u32> =
std::collections::HashMap::new();
tracing::info!("Starting multi-user heartbeat loop");
loop {
if let Some(fire_at) = config.fire_at {
let sleep_dur = duration_until_next_fire(fire_at, config.resolved_tz());
tokio::time::sleep(sleep_dur).await;
} else if let Some(ref mut iv) = tick_interval {
iv.tick().await;
}
if config.is_quiet_hours() {
continue;
}
// Get distinct user_ids from routines
let user_ids = match store.list_all_routines().await {
Ok(routines) => {
let mut ids: Vec<String> = routines
.iter()
.map(|r| r.user_id.clone())
.collect::<std::collections::HashSet<_>>()
.into_iter()
.collect();
ids.sort();
ids
}
Err(e) => {
tracing::error!("Multi-user heartbeat: failed to list routines: {}", e);
continue;
}
};
// Run user heartbeats concurrently so one slow LLM call doesn't
// block others. Cap concurrency to avoid flooding the LLM provider.
const MAX_CONCURRENT_HEARTBEATS: usize = 8;
let mut join_set = tokio::task::JoinSet::new();
for user_id in &user_ids {
// Skip users that have exceeded max_failures
let failures = user_failures.get(user_id).copied().unwrap_or(0);
if failures >= config.max_failures {
continue;
}
let workspace = Arc::new(Workspace::new_with_db(user_id, Arc::clone(store.db())));
// Run memory hygiene per user (same as single-user heartbeat).
let hygiene_ws = Arc::clone(&workspace);
let hygiene_cfg = hygiene_config.clone();
let hygiene_user = user_id.clone();
tokio::spawn(async move {
let report =
crate::workspace::hygiene::run_if_due(&hygiene_ws, &hygiene_cfg).await;
if report.had_work() {
tracing::info!(
user_id = hygiene_user,
daily_logs_deleted = report.daily_logs_deleted,
conversation_docs_deleted = report.conversation_docs_deleted,
"multi-user heartbeat: memory hygiene deleted stale documents"
);
}
});
// Drain completed tasks to stay within the concurrency cap.
while join_set.len() >= MAX_CONCURRENT_HEARTBEATS {
if let Some(join_result) = join_set.join_next().await {
collect_heartbeat_result(join_result, &mut user_failures, &config);
}
}
let uid = user_id.clone();
let cfg = config.clone();
let hyg = hygiene_config.clone();
let llm_clone = llm.clone();
let tx = response_tx.clone();
let admin = store.clone();
join_set.spawn(async move {
let mut runner = HeartbeatRunner::new(cfg, hyg, workspace, llm_clone);
if let Some(tx) = tx {
runner = runner.with_response_channel(tx);
}
runner = runner.with_store(admin);
let result = runner.check_heartbeat().await;
if let HeartbeatResult::NeedsAttention(msg) = &result {
runner.send_notification(msg).await;
}
(uid, result)
});
}
// Collect remaining results and update failure counts
while let Some(join_result) = join_set.join_next().await {
collect_heartbeat_result(join_result, &mut user_failures, &config);
}
}
})
}
/// Process a single JoinSet result from the multi-user heartbeat loop.
fn collect_heartbeat_result(
join_result: Result<(String, HeartbeatResult), tokio::task::JoinError>,
user_failures: &mut std::collections::HashMap<String, u32>,
config: &HeartbeatConfig,
) {
let (uid, result) = match join_result {
Ok(pair) => pair,
Err(e) => {
tracing::error!("Multi-user heartbeat task panicked: {}", e);
return;
}
};
match result {
HeartbeatResult::Ok => {
tracing::trace!(user_id = uid, "Multi-user heartbeat OK");
user_failures.remove(&uid);
}
HeartbeatResult::NeedsAttention(_) => {
tracing::info!(user_id = uid, "Multi-user heartbeat needs attention");
user_failures.remove(&uid);
}
HeartbeatResult::Skipped => {}
HeartbeatResult::Failed(err) => {
let count = user_failures.entry(uid.clone()).or_insert(0);
*count += 1;
tracing::error!(
user_id = uid,
consecutive_failures = *count,
"Multi-user heartbeat failed: {}",
err
);
if *count >= config.max_failures {
tracing::error!(
user_id = uid,
"Multi-user heartbeat disabled for user after {} consecutive failures",
count
);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -726,7 +903,7 @@ mod tests {
Arc<crate::workspace::Workspace>,
Arc<dyn crate::llm::LlmProvider>,
Option<tokio::sync::mpsc::Sender<crate::channels::OutgoingResponse>>,
Option<Arc<dyn crate::db::Database>>,
Option<AdminScope>,
) -> tokio::task::JoinHandle<()> = spawn_heartbeat;
let _ = _fn_ptr;
}
+3 -1
View File
@@ -36,7 +36,9 @@ pub(crate) use agent_loop::truncate_for_preview;
pub use agent_loop::{Agent, AgentDeps};
pub use compaction::{CompactionResult, ContextCompactor};
pub use context_monitor::{CompactionStrategy, ContextBreakdown, ContextMonitor};
pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat};
pub use heartbeat::{
HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_heartbeat, spawn_multi_user_heartbeat,
};
pub use router::{MessageIntent, Router};
pub use routine::{Routine, RoutineAction, RoutineRun, Trigger};
pub use routine_engine::{RoutineEngine, SandboxReadiness};
+90 -18
View File
@@ -18,6 +18,7 @@ use std::time::Duration;
use chrono::Utc;
use regex::Regex;
use tokio::sync::{RwLock, mpsc};
use tokio::task::JoinHandle;
use uuid::Uuid;
use crate::agent::Scheduler;
@@ -27,12 +28,12 @@ use crate::agent::routine::{
use crate::channels::{IncomingMessage, OutgoingResponse};
use crate::config::RoutineConfig;
use crate::context::{JobContext, JobState};
use crate::db::Database;
use crate::error::RoutineError;
use crate::extensions::ExtensionManager;
use crate::llm::{
ChatMessage, CompletionRequest, FinishReason, LlmProvider, ToolCall, ToolCompletionRequest,
};
use crate::tenant::AdminScope;
use crate::tools::{
ToolError, ToolRegistry, autonomous_allowed_tool_names, autonomous_unavailable_message,
prepare_tool_params,
@@ -45,6 +46,11 @@ enum EventMatcher {
System { routine: Routine },
}
struct TriggeredRoutine {
routine: Routine,
detail: String,
}
/// Distinguishes why sandbox is unavailable so error messages are accurate.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SandboxReadiness {
@@ -93,7 +99,7 @@ pub(crate) fn routine_matches_message(routine: &Routine, message: &IncomingMessa
/// The routine execution engine.
pub struct RoutineEngine {
config: RoutineConfig,
store: Arc<dyn Database>,
store: AdminScope,
llm: Arc<dyn LlmProvider>,
workspace: Arc<Workspace>,
/// Sender for notifications (routed to channel manager).
@@ -122,7 +128,7 @@ impl RoutineEngine {
#[allow(clippy::too_many_arguments)]
pub fn new(
config: RoutineConfig,
store: Arc<dyn Database>,
store: AdminScope,
llm: Arc<dyn LlmProvider>,
workspace: Arc<Workspace>,
notify_tx: mpsc::Sender<OutgoingResponse>,
@@ -202,6 +208,44 @@ impl RoutineEngine {
/// Check incoming message against event triggers. Returns number of routines fired.
pub async fn check_event_triggers(&self, message: &IncomingMessage, content: &str) -> usize {
let triggered = self.matching_event_triggers(message, content).await;
let fired = triggered.len();
for triggered in triggered {
std::mem::drop(self.spawn_fire(triggered.routine, "event", Some(triggered.detail)));
}
fired
}
/// Fire matching event-triggered routines and wait for them to complete.
///
/// Used by single-message REPL mode so the process does not exit before
/// background event-triggered routines finish.
pub async fn check_event_triggers_and_wait(
&self,
message: &IncomingMessage,
content: &str,
) -> usize {
let triggered = self.matching_event_triggers(message, content).await;
let fired = triggered.len();
let handles: Vec<JoinHandle<()>> = triggered
.into_iter()
.map(|triggered| self.spawn_fire(triggered.routine, "event", Some(triggered.detail)))
.collect();
for handle in handles {
if let Err(e) = handle.await {
tracing::warn!(error = %e, "Event-triggered routine task failed");
}
}
fired
}
async fn matching_event_triggers(
&self,
message: &IncomingMessage,
content: &str,
) -> Vec<TriggeredRoutine> {
let cache = self.event_cache.read().await;
// Early return if there are no message matchers at all.
@@ -209,10 +253,9 @@ impl RoutineEngine {
.iter()
.any(|m| matches!(m, EventMatcher::Message { .. }))
{
return 0;
return Vec::new();
}
let mut fired = 0;
let mut triggered = Vec::new();
// Collect routine IDs for batch query
let routine_ids: Vec<Uuid> = cache
@@ -224,13 +267,13 @@ impl RoutineEngine {
.collect();
if routine_ids.is_empty() {
return 0;
return Vec::new();
}
// Single batch query instead of N queries
let concurrent_counts = match self.batch_concurrent_counts(&routine_ids).await {
Some(counts) => counts,
None => return 0,
None => return Vec::new(),
};
for matcher in cache.iter() {
@@ -285,11 +328,13 @@ impl RoutineEngine {
}
let detail = truncate(content, 200);
self.spawn_fire(routine.clone(), "event", Some(detail));
fired += 1;
triggered.push(TriggeredRoutine {
routine: routine.clone(),
detail,
});
}
fired
triggered
}
/// Emit a structured event to system-event routines.
@@ -737,12 +782,22 @@ impl RoutineEngine {
});
}
// Per-user workspace (same pattern as spawn_fire).
let routine_workspace = if routine.user_id == self.workspace.user_id() {
self.workspace.clone()
} else {
Arc::new(Workspace::new_with_db(
&routine.user_id,
Arc::clone(self.store.db()),
))
};
// Execute inline for manual triggers (caller wants to wait)
let engine = EngineContext {
config: self.config.clone(),
store: self.store.clone(),
llm: self.llm.clone(),
workspace: self.workspace.clone(),
workspace: routine_workspace,
notify_tx: self.notify_tx.clone(),
running_count: self.running_count.clone(),
scheduler: self.scheduler.clone(),
@@ -845,7 +900,12 @@ impl RoutineEngine {
}
/// Spawn a fire in a background task.
fn spawn_fire(&self, routine: Routine, trigger_type: &str, trigger_detail: Option<String>) {
fn spawn_fire(
&self,
routine: Routine,
trigger_type: &str,
trigger_detail: Option<String>,
) -> JoinHandle<()> {
let run = RoutineRun {
id: Uuid::new_v4(),
routine_id: routine.id,
@@ -860,11 +920,23 @@ impl RoutineEngine {
created_at: Utc::now(),
};
// Use per-user workspace so each routine executes in the correct
// user's context. Fall back to the engine-wide workspace when the
// routine belongs to the same user (avoids unnecessary allocation).
let routine_workspace = if routine.user_id == self.workspace.user_id() {
self.workspace.clone()
} else {
Arc::new(Workspace::new_with_db(
&routine.user_id,
Arc::clone(self.store.db()),
))
};
let engine = EngineContext {
config: self.config.clone(),
store: self.store.clone(),
llm: self.llm.clone(),
workspace: self.workspace.clone(),
workspace: routine_workspace,
notify_tx: self.notify_tx.clone(),
running_count: self.running_count.clone(),
scheduler: self.scheduler.clone(),
@@ -882,7 +954,7 @@ impl RoutineEngine {
return;
}
execute_routine(engine, routine, run).await;
});
})
}
fn check_cooldown(&self, routine: &Routine) -> bool {
@@ -917,7 +989,7 @@ impl RoutineEngine {
/// an active state (Pending/InProgress/Stuck). Maps the final `JobState` to
/// a `RunStatus` for the routine run.
struct FullJobWatcher {
store: Arc<dyn Database>,
store: AdminScope,
job_id: Uuid,
routine_name: String,
}
@@ -928,7 +1000,7 @@ impl FullJobWatcher {
/// Safety ceiling: 24 hours, derived from POLL_INTERVAL.
const MAX_POLLS: u32 = (24 * 60 * 60) / Self::POLL_INTERVAL.as_secs() as u32;
fn new(store: Arc<dyn Database>, job_id: Uuid, routine_name: String) -> Self {
fn new(store: AdminScope, job_id: Uuid, routine_name: String) -> Self {
Self {
store,
job_id,
@@ -1000,7 +1072,7 @@ impl FullJobWatcher {
/// Shared context passed to the execution function.
struct EngineContext {
config: RoutineConfig,
store: Arc<dyn Database>,
store: AdminScope,
llm: Arc<dyn LlmProvider>,
workspace: Arc<Workspace>,
notify_tx: mpsc::Sender<OutgoingResponse>,
+7 -3
View File
@@ -11,12 +11,12 @@ use uuid::Uuid;
use crate::agent::task::{Task, TaskContext, TaskOutput};
use crate::config::AgentConfig;
use crate::context::{ContextManager, JobContext, JobState};
use crate::db::Database;
use crate::error::{Error, JobError};
use crate::extensions::ExtensionManager;
use crate::hooks::HookRegistry;
use crate::llm::LlmProvider;
use crate::safety::SafetyLayer;
use crate::tenant::AdminScope;
use crate::tools::{
ApprovalContext, ToolRegistry, autonomous_allowed_tool_names, autonomous_unavailable_error,
prepare_tool_params,
@@ -52,7 +52,7 @@ struct ScheduledSubtask {
pub struct SchedulerDeps {
pub tools: Arc<ToolRegistry>,
pub extension_manager: Option<Arc<ExtensionManager>>,
pub store: Option<Arc<dyn Database>>,
pub store: Option<AdminScope>,
pub hooks: Arc<HookRegistry>,
}
@@ -64,7 +64,7 @@ pub struct Scheduler {
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
extension_manager: Option<Arc<ExtensionManager>>,
store: Option<Arc<dyn Database>>,
store: Option<AdminScope>,
hooks: Arc<HookRegistry>,
/// SSE manager for live job event streaming.
sse_tx: Option<Arc<crate::channels::web::sse::SseManager>>,
@@ -780,10 +780,14 @@ mod tests {
allow_local_tools: true,
max_cost_per_day_cents: None,
max_actions_per_hour: None,
max_cost_per_user_per_day_cents: None,
max_tool_iterations: 10,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
max_tokens_per_job,
multi_tenant: false,
max_llm_concurrent_per_user: None,
max_jobs_concurrent_per_user: None,
};
let cm = Arc::new(ContextManager::new(5));
let llm: Arc<dyn LlmProvider> = Arc::new(StubLlm);
+5 -5
View File
@@ -8,8 +8,8 @@ use chrono::{DateTime, Utc};
use uuid::Uuid;
use crate::context::{ContextManager, JobState};
use crate::db::Database;
use crate::error::RepairError;
use crate::tenant::AdminScope;
use crate::tools::{BuildRequirement, Language, SoftwareBuilder, SoftwareType, ToolRegistry};
/// A job that has been detected as stuck.
@@ -69,7 +69,7 @@ pub struct DefaultSelfRepair {
/// Jobs in `InProgress` longer than this are treated as stuck.
stuck_threshold: Duration,
max_repair_attempts: u32,
store: Option<Arc<dyn Database>>,
store: Option<AdminScope>,
builder: Option<Arc<dyn SoftwareBuilder>>,
tools: Option<Arc<ToolRegistry>>,
}
@@ -91,8 +91,8 @@ impl DefaultSelfRepair {
}
}
/// Add a Store for tool failure tracking.
pub fn with_store(mut self, store: Arc<dyn Database>) -> Self {
/// Add an admin-scoped store for tool failure tracking.
pub fn with_store(mut self, store: AdminScope) -> Self {
self.store = Some(store);
self
}
@@ -806,7 +806,7 @@ mod tests {
// Create self-repair with zero threshold (detect immediately),
// wired with store, builder, and tools.
let repair = DefaultSelfRepair::new(Arc::clone(&cm), Duration::from_secs(0), 3)
.with_store(Arc::clone(&db))
.with_store(crate::tenant::AdminScope::new(Arc::clone(&db)))
.with_builder(
Arc::clone(&builder) as Arc<dyn crate::tools::SoftwareBuilder>,
tools,
+192 -1
View File
@@ -449,6 +449,7 @@ impl Thread {
id: call_id.clone(),
name: tc.name.clone(),
arguments: tc.parameters.clone(),
reasoning: None,
})
.collect();
@@ -522,7 +523,12 @@ impl Thread {
&& let Some(ref tcs) = assistant_msg.tool_calls
{
for tc in tcs {
turn.record_tool_call(&tc.name, tc.arguments.clone());
turn.record_tool_call_with_reasoning(
&tc.name,
tc.arguments.clone(),
tc.reasoning.clone(),
Some(tc.id.clone()),
);
}
}
@@ -602,6 +608,10 @@ pub struct Turn {
pub completed_at: Option<DateTime<Utc>>,
/// Error message (if failed).
pub error: Option<String>,
/// Agent's reasoning narrative for this turn.
/// Cleaned via `clean_response` and sanitized through `SafetyLayer` before storage.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub narrative: Option<String>,
/// Transient image content parts for multimodal LLM input.
/// Not serialized — images are only needed for the current LLM call.
/// The text description in `user_input` persists for compaction/context.
@@ -621,6 +631,7 @@ impl Turn {
started_at: Utc::now(),
completed_at: None,
error: None,
narrative: None,
image_content_parts: Vec::new(),
}
}
@@ -656,6 +667,26 @@ impl Turn {
parameters: params,
result: None,
error: None,
rationale: None,
tool_call_id: None,
});
}
/// Record a tool call with reasoning context.
pub fn record_tool_call_with_reasoning(
&mut self,
name: impl Into<String>,
params: serde_json::Value,
rationale: Option<String>,
tool_call_id: Option<String>,
) {
self.tool_calls.push(TurnToolCall {
name: name.into(),
parameters: params,
result: None,
error: None,
rationale,
tool_call_id,
});
}
@@ -672,6 +703,60 @@ impl Turn {
call.error = Some(error.into());
}
}
/// Record a tool result by tool_call_id, with fallback to first pending call.
pub fn record_tool_result_for(&mut self, tool_call_id: &str, result: serde_json::Value) {
if let Some(call) = self
.tool_calls
.iter_mut()
.find(|c| c.tool_call_id.as_deref() == Some(tool_call_id))
{
call.result = Some(result);
} else if let Some(call) = self
.tool_calls
.iter_mut()
.find(|c| c.result.is_none() && c.error.is_none())
{
tracing::debug!(
tool_call_id = %tool_call_id,
fallback_tool = %call.name,
"tool_call_id not found, falling back to first pending call"
);
call.result = Some(result);
} else {
tracing::warn!(
tool_call_id = %tool_call_id,
"Tool result dropped: no matching or pending tool call"
);
}
}
/// Record a tool error by tool_call_id, with fallback to first pending call.
pub fn record_tool_error_for(&mut self, tool_call_id: &str, error: impl Into<String>) {
if let Some(call) = self
.tool_calls
.iter_mut()
.find(|c| c.tool_call_id.as_deref() == Some(tool_call_id))
{
call.error = Some(error.into());
} else if let Some(call) = self
.tool_calls
.iter_mut()
.find(|c| c.result.is_none() && c.error.is_none())
{
tracing::debug!(
tool_call_id = %tool_call_id,
fallback_tool = %call.name,
"tool_call_id not found, falling back to first pending call"
);
call.error = Some(error.into());
} else {
tracing::warn!(
tool_call_id = %tool_call_id,
"Tool error dropped: no matching or pending tool call"
);
}
}
}
/// Record of a tool call made during a turn.
@@ -685,6 +770,12 @@ pub struct TurnToolCall {
pub result: Option<serde_json::Value>,
/// Error from the tool (if failed).
pub error: Option<String>,
/// Agent's reasoning for choosing this tool.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub rationale: Option<String>,
/// The tool_call_id from the LLM, for identity-based result matching.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
}
#[cfg(test)]
@@ -1309,6 +1400,7 @@ mod tests {
id: "call_0".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"q": "test"}),
reasoning: None,
};
let messages = vec![
ChatMessage::user("Find test"),
@@ -1339,6 +1431,7 @@ mod tests {
id: "call_0".to_string(),
name: "http".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
};
let messages = vec![
ChatMessage::user("Fetch URL"),
@@ -1404,11 +1497,13 @@ mod tests {
id: "call_a".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"q": "data"}),
reasoning: None,
};
let tc2 = ToolCall {
id: "call_b".to_string(),
name: "write".to_string(),
arguments: serde_json::json!({"path": "out.txt"}),
reasoning: None,
};
let messages = vec![
ChatMessage::user("Find and save"),
@@ -1620,4 +1715,100 @@ mod tests {
let merged = thread.drain_pending_messages().unwrap();
assert_eq!(merged, "failed batch\nnew msg");
}
#[test]
fn test_record_tool_result_for_by_id() {
let mut turn = Turn::new(0, "test");
turn.record_tool_call_with_reasoning(
"tool_a",
serde_json::json!({}),
None,
Some("id_a".into()),
);
turn.record_tool_call_with_reasoning(
"tool_b",
serde_json::json!({}),
None,
Some("id_b".into()),
);
// Record result for second tool by ID
turn.record_tool_result_for("id_b", serde_json::json!("result_b"));
assert!(turn.tool_calls[0].result.is_none());
assert_eq!(
turn.tool_calls[1].result.as_ref().unwrap(),
&serde_json::json!("result_b")
);
}
#[test]
fn test_record_tool_error_for_by_id() {
let mut turn = Turn::new(0, "test");
turn.record_tool_call_with_reasoning(
"tool_a",
serde_json::json!({}),
None,
Some("id_a".into()),
);
turn.record_tool_call_with_reasoning(
"tool_b",
serde_json::json!({}),
None,
Some("id_b".into()),
);
turn.record_tool_error_for("id_a", "failed");
assert_eq!(turn.tool_calls[0].error.as_deref(), Some("failed"));
assert!(turn.tool_calls[1].error.is_none());
}
#[test]
fn test_record_tool_result_for_fallback_to_pending() {
let mut turn = Turn::new(0, "test");
turn.record_tool_call_with_reasoning(
"tool_a",
serde_json::json!({}),
None,
Some("id_a".into()),
);
turn.record_tool_call_with_reasoning(
"tool_b",
serde_json::json!({}),
None,
Some("id_b".into()),
);
// First tool already has a result
turn.tool_calls[0].result = Some(serde_json::json!("done"));
// Unknown ID should fall back to first pending (tool_b)
turn.record_tool_result_for("unknown_id", serde_json::json!("fallback"));
assert_eq!(
turn.tool_calls[0].result.as_ref().unwrap(),
&serde_json::json!("done")
);
assert_eq!(
turn.tool_calls[1].result.as_ref().unwrap(),
&serde_json::json!("fallback")
);
}
#[test]
fn test_record_tool_result_for_no_pending_is_noop() {
let mut turn = Turn::new(0, "test");
turn.record_tool_call_with_reasoning(
"tool_a",
serde_json::json!({}),
None,
Some("id_a".into()),
);
turn.tool_calls[0].result = Some(serde_json::json!("done"));
// No pending calls, unknown ID — should be a no-op
turn.record_tool_result_for("unknown_id", serde_json::json!("lost"));
assert_eq!(
turn.tool_calls[0].result.as_ref().unwrap(),
&serde_json::json!("done")
);
}
}
+11
View File
@@ -92,6 +92,17 @@ impl SubmissionParser {
args: vec![],
};
}
if lower == "/reasoning" || lower.starts_with("/reasoning ") {
let args: Vec<String> = trimmed
.split_whitespace()
.skip(1)
.map(|s| s.to_string())
.collect();
return Submission::SystemCommand {
command: "reasoning".to_string(),
args,
};
}
if lower == "/restart" {
tracing::debug!("[SubmissionParser::parse] Recognized /restart command");
return Submission::SystemCommand {
+68 -14
View File
@@ -175,6 +175,7 @@ impl Agent {
pub(super) async fn process_user_input(
&self,
message: &IncomingMessage,
tenant: crate::tenant::TenantCtx,
session: Arc<Mutex<Session>>,
thread_id: Uuid,
content: &str,
@@ -351,7 +352,7 @@ impl Agent {
if let Some(intent) = self.router.route_command(&temp_message) {
// Explicit command like /status, /job, /list - handle directly
return self.handle_job_or_command(intent, message).await;
return self.handle_job_or_command(intent, message, &tenant).await;
}
// Natural language goes through the agentic loop
@@ -462,7 +463,7 @@ impl Agent {
// Run the agentic tool execution loop
let result = self
.run_agentic_loop(message, session.clone(), thread_id, turn_messages)
.run_agentic_loop(message, tenant, session.clone(), thread_id, turn_messages)
.await;
// Re-acquire lock and check if interrupted
@@ -513,10 +514,10 @@ impl Agent {
};
thread.complete_turn(&response);
let (turn_number, tool_calls) = thread
let (turn_number, tool_calls, narrative) = thread
.turns
.last()
.map(|t| (t.turn_number, t.tool_calls.clone()))
.map(|t| (t.turn_number, t.tool_calls.clone(), t.narrative.clone()))
.unwrap_or_default();
let _ = self
.channels
@@ -534,6 +535,7 @@ impl Agent {
&message.user_id,
turn_number,
&tool_calls,
narrative.as_deref(),
)
.await;
self.persist_assistant_response(
@@ -725,7 +727,9 @@ impl Agent {
///
/// Stored between the user and assistant messages so that
/// `build_turns_from_db_messages` can reconstruct the tool call history.
/// Content is a JSON array of tool call summaries.
/// Content is a JSON object: `{ "calls": [...], "narrative": "..." }`.
/// The `calls` array contains tool call summaries with optional `rationale`
/// and `tool_call_id` fields. Legacy rows may be plain JSON arrays.
pub(super) async fn persist_tool_calls(
&self,
thread_id: Uuid,
@@ -733,6 +737,7 @@ impl Agent {
user_id: &str,
turn_number: usize,
tool_calls: &[crate::agent::session::TurnToolCall],
narrative: Option<&str>,
) {
if tool_calls.is_empty() {
return;
@@ -767,11 +772,30 @@ impl Agent {
if let Some(ref error) = tc.error {
obj["error"] = serde_json::Value::String(truncate_preview(error, 200));
}
if let Some(ref rationale) = tc.rationale {
obj["rationale"] = serde_json::Value::String(truncate_preview(rationale, 500));
}
if let Some(ref tool_call_id) = tc.tool_call_id {
obj["tool_call_id"] =
serde_json::Value::String(truncate_preview(tool_call_id, 128));
}
obj
})
.collect();
let content = match serde_json::to_string(&summaries) {
// Wrap in an object with optional narrative so it can be reconstructed.
// safety: no byte-index slicing here; comment describes JSON shape
let wrapper = if let Some(n) = narrative {
serde_json::json!({
"narrative": truncate_preview(n, 1000),
"calls": summaries,
})
} else {
serde_json::json!({
"calls": summaries,
})
};
let content = match serde_json::to_string(&wrapper) {
Ok(c) => c,
Err(e) => {
tracing::warn!("Failed to serialize tool calls: {}", e);
@@ -1104,9 +1128,12 @@ impl Agent {
&& let Some(turn) = thread.last_turn_mut()
{
if is_tool_error {
turn.record_tool_error(result_content.clone());
turn.record_tool_error_for(&pending.tool_call_id, result_content.clone());
} else {
turn.record_tool_result(serde_json::json!(result_content));
turn.record_tool_result_for(
&pending.tool_call_id,
serde_json::json!(result_content),
);
}
}
}
@@ -1358,9 +1385,12 @@ impl Agent {
&& let Some(turn) = thread.last_turn_mut()
{
if is_deferred_error {
turn.record_tool_error(deferred_content.clone());
turn.record_tool_error_for(&tc.id, deferred_content.clone());
} else {
turn.record_tool_result(serde_json::json!(deferred_content));
turn.record_tool_result_for(
&tc.id,
serde_json::json!(deferred_content),
);
}
}
}
@@ -1444,7 +1474,13 @@ impl Agent {
// Continue the agentic loop (a tool was already executed this turn)
let result = self
.run_agentic_loop(message, session.clone(), thread_id, context_messages)
.run_agentic_loop(
message,
self.tenant_ctx(&message.user_id).await,
session.clone(),
thread_id,
context_messages,
)
.await;
// Handle the result
@@ -1459,10 +1495,10 @@ impl Agent {
let (response, suggestions) =
crate::agent::dispatcher::extract_suggestions(&response);
thread.complete_turn(&response);
let (turn_number, tool_calls) = thread
let (turn_number, tool_calls, narrative) = thread
.turns
.last()
.map(|t| (t.turn_number, t.tool_calls.clone()))
.map(|t| (t.turn_number, t.tool_calls.clone(), t.narrative.clone()))
.unwrap_or_default();
// User message already persisted at turn start; save tool calls then assistant response
self.persist_tool_calls(
@@ -1471,6 +1507,7 @@ impl Agent {
&message.user_id,
turn_number,
&tool_calls,
narrative.as_deref(),
)
.await;
self.persist_assistant_response(
@@ -1816,7 +1853,20 @@ fn rebuild_chat_messages_from_db(
"assistant" => result.push(ChatMessage::assistant(&msg.content)),
"tool_calls" => {
// Try to parse the enriched JSON and rebuild tool messages.
if let Ok(calls) = serde_json::from_str::<Vec<serde_json::Value>>(&msg.content) {
// Supports two formats:
// - Old: plain JSON array of tool call summaries
// - New: wrapped object { "calls": [...], "narrative": "..." }
let calls: Vec<serde_json::Value> =
match serde_json::from_str::<serde_json::Value>(&msg.content) {
Ok(serde_json::Value::Array(arr)) => arr,
Ok(serde_json::Value::Object(obj)) => obj
.get("calls")
.and_then(|v| v.as_array())
.cloned()
.unwrap_or_default(),
_ => Vec::new(),
};
{
if calls.is_empty() {
continue;
}
@@ -1839,6 +1889,10 @@ fn rebuild_chat_messages_from_db(
.get("parameters")
.cloned()
.unwrap_or(serde_json::json!({})),
reasoning: c
.get("rationale")
.and_then(|v| v.as_str())
.map(String::from),
})
.collect();
+1
View File
@@ -897,6 +897,7 @@ impl AppBuilder {
crate::agent::cost_guard::CostGuardConfig {
max_cost_per_day_cents: self.config.agent.max_cost_per_day_cents,
max_actions_per_hour: self.config.agent.max_actions_per_hour,
max_cost_per_user_per_day_cents: self.config.agent.max_cost_per_user_per_day_cents,
},
));
+16
View File
@@ -265,6 +265,15 @@ impl OutgoingResponse {
}
}
/// A single tool decision within a reasoning update.
#[derive(Debug, Clone)]
pub struct ToolDecision {
/// Tool name.
pub tool_name: String,
/// Agent's reasoning for choosing this tool.
pub rationale: String,
}
/// Status update types for showing agent activity.
#[derive(Debug, Clone)]
pub enum StatusUpdate {
@@ -333,6 +342,13 @@ pub enum StatusUpdate {
},
/// Suggested follow-up messages for the user.
Suggestions { suggestions: Vec<String> },
/// Agent reasoning update (why it chose specific tools).
ReasoningUpdate {
/// Human-readable summary of the agent's decision.
narrative: String,
/// Per-tool decisions.
decisions: Vec<ToolDecision>,
},
/// Per-turn token usage and cost summary (shown as subtle metadata).
TurnCost {
input_tokens: u64,
+1 -1
View File
@@ -39,7 +39,7 @@ mod webhook_server;
pub use channel::{
AttachmentKind, Channel, ChannelSecretUpdater, IncomingAttachment, IncomingMessage,
MessageStream, OutgoingResponse, StatusUpdate, routing_target_from_metadata,
MessageStream, OutgoingResponse, StatusUpdate, ToolDecision, routing_target_from_metadata,
};
pub use http::{HttpChannel, HttpChannelState};
pub use manager::ChannelManager;
+65 -9
View File
@@ -75,6 +75,7 @@ const SLASH_COMMANDS: &[&str] = &[
"/suggest",
"/thread",
"/resume",
"/reasoning",
];
/// Rustyline helper for slash-command tab completion.
@@ -430,6 +431,18 @@ impl ReplChannel {
let _ = execute!(stderr, terminal::Clear(terminal::ClearType::FromCursorDown));
}
}
async fn finish_single_message_turn(&self) {
if self.single_message.is_none() {
return;
}
let tx = self.msg_tx.lock().ok().and_then(|mut guard| guard.take());
if let Some(tx) = tx {
let msg = IncomingMessage::new("repl", &self.user_id, "/quit");
let _ = tx.send(msg).await;
}
}
}
impl Default for ReplChannel {
@@ -479,7 +492,9 @@ impl Channel for ReplChannel {
async fn start(&self) -> Result<MessageStream, ChannelError> {
let (tx, rx) = mpsc::channel(32);
// Store tx so send_status can inject approval responses directly
// Approval prompts inject responses back through this sender.
// In single-message mode we keep it until the turn finishes, then
// drop it after enqueuing /quit so the receiver stream can close.
if let Ok(mut guard) = self.msg_tx.lock() {
*guard = Some(tx.clone());
}
@@ -495,11 +510,10 @@ impl Channel for ReplChannel {
// Single message mode: send it and return
if let Some(msg) = single_message {
let incoming = IncomingMessage::new("repl", &user_id, &msg).with_timezone(&sys_tz);
let incoming = IncomingMessage::new("repl", &user_id, &msg)
.with_metadata(serde_json::json!({ "single_message_mode": true }))
.with_timezone(&sys_tz);
let _ = tx.blocking_send(incoming);
// Ensure the agent exits after handling exactly one turn in -m mode,
// even when other channels (gateway/http) are enabled.
let _ = tx.blocking_send(IncomingMessage::new("repl", &user_id, "/quit"));
return;
}
@@ -662,6 +676,7 @@ impl Channel for ReplChannel {
println!();
println!();
self.stdin_locked.store(false, Ordering::Relaxed);
self.finish_single_message_turn().await;
return Ok(());
}
@@ -680,6 +695,7 @@ impl Channel for ReplChannel {
println!();
// Unlock stdin so readline can resume
self.stdin_locked.store(false, Ordering::Relaxed);
self.finish_single_message_turn().await;
Ok(())
}
@@ -779,6 +795,7 @@ impl Channel for ReplChannel {
let msg_tx = Arc::clone(&self.msg_tx);
let user_id = self.user_id.clone();
let lock_flag = Arc::clone(&self.stdin_locked);
let single_message_mode = self.single_message.is_some();
tokio::task::spawn_blocking(move || {
let action = run_approval_selector(allow_always).unwrap_or("n");
// Unlock stdin so readline can resume after approval
@@ -787,7 +804,12 @@ impl Channel for ReplChannel {
return;
};
if let Some(tx) = guard.as_ref() {
let msg = IncomingMessage::new("repl", &user_id, action);
let msg = if single_message_mode {
IncomingMessage::new("repl", &user_id, action)
.with_metadata(serde_json::json!({ "single_message_mode": true }))
} else {
IncomingMessage::new("repl", &user_id, action)
};
let _ = tx.blocking_send(msg);
}
});
@@ -841,6 +863,19 @@ impl Channel for ReplChannel {
StatusUpdate::Suggestions { .. } => {
// Suggestions are only rendered by the web gateway
}
StatusUpdate::ReasoningUpdate {
narrative,
decisions,
} => {
if !narrative.is_empty() {
let display = truncate_for_preview(&narrative, CLI_STATUS_MAX);
eprintln!(" \x1b[94m\u{25B6} {display}\x1b[0m");
}
for d in &decisions {
let display = truncate_for_preview(&d.rationale, CLI_STATUS_MAX);
eprintln!(" \x1b[90m\u{2192} {}: {display}\x1b[0m", d.tool_name);
}
}
StatusUpdate::TurnCost { .. } => {
// Cost display is handled by the TUI channel
}
@@ -875,6 +910,7 @@ impl Channel for ReplChannel {
#[cfg(test)]
mod tests {
use futures::StreamExt;
use tokio::time::{Duration, timeout};
use super::*;
@@ -883,16 +919,36 @@ mod tests {
let repl = ReplChannel::with_message("hi".to_string());
let mut stream = repl.start().await.expect("repl start should succeed");
let first = stream.next().await.expect("first message missing");
let first = timeout(Duration::from_secs(1), stream.next())
.await
.expect("timed out waiting for first message")
.expect("first message missing");
assert_eq!(first.channel, "repl");
assert_eq!(first.content, "hi");
let second = stream.next().await.expect("quit message missing");
assert!(
timeout(Duration::from_millis(100), stream.next())
.await
.is_err(),
"single-message mode should wait for the turn to finish before quitting"
);
repl.respond(&first, OutgoingResponse::text("done"))
.await
.expect("respond should succeed");
let second = timeout(Duration::from_secs(1), stream.next())
.await
.expect("timed out waiting for quit message")
.expect("quit message missing");
assert_eq!(second.channel, "repl");
assert_eq!(second.content, "/quit");
assert!(
stream.next().await.is_none(),
timeout(Duration::from_secs(1), stream.next())
.await
.expect("timed out waiting for stream to close")
.is_none(),
"stream should end after /quit"
);
}
+14
View File
@@ -3061,6 +3061,20 @@ fn status_to_wit(
},
// Suggestions and turn cost are web-gateway-only; skip for WASM channels
StatusUpdate::Suggestions { .. } | StatusUpdate::TurnCost { .. } => return None,
StatusUpdate::ReasoningUpdate {
narrative,
decisions,
} => {
let mut msg = narrative.clone();
for d in decisions {
msg.push_str(&format!("\n{}: {}", d.tool_name, d.rationale));
}
wit_channel::StatusUpdate {
status: wit_channel::StatusType::Status,
message: msg,
metadata_json,
}
}
})
}
+2
View File
@@ -398,8 +398,10 @@ pub async fn chat_history_handler(
truncate_preview(&s, 500)
}),
error: tc.error.clone(),
rationale: tc.rationale.clone(),
})
.collect(),
narrative: t.narrative.clone(),
})
.collect();
+30 -3
View File
@@ -54,10 +54,37 @@ fn validate_webhook_secret(
///
/// This endpoint is **public** (no gateway auth token required) but protected
/// by the per-routine webhook secret sent via the `X-Webhook-Secret` header.
///
/// **Single-user/backward-compatible**: looks up routines by path across all
/// users. For multi-tenant isolation, use the user-scoped endpoint at
/// `/api/webhooks/u/{user_id}/{path}` instead.
pub async fn webhook_trigger_handler(
State(state): State<Arc<GatewayState>>,
Path(path): Path<String>,
headers: HeaderMap,
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
fire_webhook_inner(state, &path, None, &headers).await
}
/// Handle incoming webhook POST to `/api/webhooks/u/{user_id}/{path}`.
///
/// User-scoped variant for multi-tenant deployments. The `user_id` in the URL
/// restricts the routine lookup to that user only, preventing cross-user
/// webhook triggering even when paths collide.
pub async fn webhook_trigger_user_scoped_handler(
State(state): State<Arc<GatewayState>>,
Path((user_id, path)): Path<(String, String)>,
headers: HeaderMap,
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
fire_webhook_inner(state, &path, Some(&user_id), &headers).await
}
/// Shared webhook logic for both scoped and unscoped endpoints.
async fn fire_webhook_inner(
state: Arc<GatewayState>,
path: &str,
user_id: Option<&str>,
headers: &HeaderMap,
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
// Rate limit check
if !state.webhook_rate_limiter.check() {
@@ -72,9 +99,9 @@ pub async fn webhook_trigger_handler(
"Database not available".to_string(),
))?;
// Targeted query instead of loading all routines
// Targeted query — when user_id is provided, restrict to that user's routines
let routine = store
.get_webhook_routine_by_path(&path)
.get_webhook_routine_by_path(path, user_id)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
.ok_or((
@@ -99,7 +126,7 @@ pub async fn webhook_trigger_handler(
))?
};
let run_id = engine.fire_webhook(routine.id, &path).await.map_err(|e| {
let run_id = engine.fire_webhook(routine.id, path).await.map_err(|e| {
let status = match &e {
crate::error::RoutineError::NotFound { .. } => StatusCode::NOT_FOUND,
crate::error::RoutineError::Disabled { .. }
+14
View File
@@ -501,6 +501,20 @@ impl Channel for GatewayChannel {
},
StatusUpdate::Suggestions { suggestions } => AppEvent::Suggestions {
suggestions,
thread_id: thread_id.clone(),
},
StatusUpdate::ReasoningUpdate {
narrative,
decisions,
} => AppEvent::ReasoningUpdate {
narrative,
decisions: decisions
.into_iter()
.map(|d| crate::channels::web::types::ToolDecisionDto {
tool_name: d.tool_name,
rationale: d.rationale,
})
.collect(),
thread_id,
},
StatusUpdate::TurnCost {
+2
View File
@@ -231,6 +231,7 @@ pub fn convert_messages(messages: &[OpenAiMessage]) -> Result<Vec<ChatMessage>,
name: tc.function.name.clone(),
arguments: serde_json::from_str(&tc.function.arguments)
.unwrap_or(serde_json::Value::Object(Default::default())),
reasoning: None,
})
.collect();
Ok(ChatMessage::assistant_with_tool_calls(
@@ -954,6 +955,7 @@ mod tests {
id: "call_abc".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"query": "rust"}),
reasoning: None,
}];
let converted = convert_tool_calls_to_openai(&calls);
+7
View File
@@ -421,6 +421,11 @@ pub async fn start_server(
.route(
"/api/webhooks/{path}",
post(crate::channels::web::handlers::webhooks::webhook_trigger_handler),
)
// User-scoped webhook endpoint for multi-tenant isolation
.route(
"/api/webhooks/u/{user_id}/{path}",
post(crate::channels::web::handlers::webhooks::webhook_trigger_user_scoped_handler),
);
// Protected routes (require auth)
@@ -1750,8 +1755,10 @@ async fn chat_history_handler(
truncate_preview(&s, 500)
}),
error: tc.error.clone(),
rationale: tc.rationale.clone(),
})
.collect(),
narrative: t.narrative.clone(),
})
.collect();
+7 -1
View File
@@ -63,6 +63,9 @@ pub struct TurnInfo {
pub started_at: String,
pub completed_at: Option<String>,
pub tool_calls: Vec<ToolCallInfo>,
/// Agent's reasoning narrative for this turn.
#[serde(skip_serializing_if = "Option::is_none")]
pub narrative: Option<String>,
}
#[derive(Debug, Serialize)]
@@ -74,6 +77,9 @@ pub struct ToolCallInfo {
pub result_preview: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
/// Agent's reasoning for choosing this tool.
#[serde(skip_serializing_if = "Option::is_none")]
pub rationale: Option<String>,
}
#[derive(Debug, Serialize)]
@@ -116,7 +122,7 @@ pub struct ApprovalRequest {
// --- App Event (re-exported from ironclaw_common) ---
pub use ironclaw_common::AppEvent;
pub use ironclaw_common::{AppEvent, ToolDecisionDto};
// --- Memory ---
+87 -12
View File
@@ -4,6 +4,21 @@ use crate::channels::web::types::{ToolCallInfo, TurnInfo};
pub use ironclaw_common::truncate_preview;
/// Parse tool call summary JSON objects into `ToolCallInfo` structs.
fn parse_tool_call_infos(calls: &[serde_json::Value]) -> Vec<ToolCallInfo> {
calls
.iter()
.map(|c| ToolCallInfo {
name: c["name"].as_str().unwrap_or("unknown").to_string(),
has_result: c.get("result_preview").is_some_and(|v| !v.is_null()),
has_error: c.get("error").is_some_and(|v| !v.is_null()),
result_preview: c["result_preview"].as_str().map(String::from),
error: c["error"].as_str().map(String::from),
rationale: c["rationale"].as_str().map(String::from),
})
.collect()
}
/// Build TurnInfo pairs from flat DB messages (user/tool_calls/assistant triples).
///
/// Handles three message patterns:
@@ -27,6 +42,7 @@ pub fn build_turns_from_db_messages(
started_at: msg.created_at.to_rfc3339(),
completed_at: None,
tool_calls: Vec::new(),
narrative: None,
};
// Check if next message is a tool_calls record
@@ -34,18 +50,28 @@ pub fn build_turns_from_db_messages(
&& next.role == "tool_calls"
{
let tc_msg = iter.next().expect("peeked");
match serde_json::from_str::<Vec<serde_json::Value>>(&tc_msg.content) {
Ok(calls) => {
turn.tool_calls = calls
.iter()
.map(|c| ToolCallInfo {
name: c["name"].as_str().unwrap_or("unknown").to_string(),
has_result: c.get("result_preview").is_some(),
has_error: c.get("error").is_some(),
result_preview: c["result_preview"].as_str().map(String::from),
error: c["error"].as_str().map(String::from),
})
.collect();
// Parse tool_calls JSON — supports two formats:
// safety: no byte-index slicing; comment describes JSON shape
match serde_json::from_str::<serde_json::Value>(&tc_msg.content) {
Ok(serde_json::Value::Array(calls)) => {
// Old format: plain array
turn.tool_calls = parse_tool_call_infos(&calls);
}
Ok(serde_json::Value::Object(obj)) => {
// New wrapped format with narrative
turn.narrative = obj
.get("narrative")
.and_then(|v| v.as_str())
.map(String::from);
if let Some(serde_json::Value::Array(calls)) = obj.get("calls") {
turn.tool_calls = parse_tool_call_infos(calls);
}
}
Ok(_) => {
tracing::warn!(
message_id = %tc_msg.id,
"Unexpected tool_calls JSON shape in DB, skipping"
);
}
Err(e) => {
tracing::warn!(
@@ -83,6 +109,7 @@ pub fn build_turns_from_db_messages(
started_at: msg.created_at.to_rfc3339(),
completed_at: Some(msg.created_at.to_rfc3339()),
tool_calls: Vec::new(),
narrative: None,
});
turn_number += 1;
}
@@ -201,4 +228,52 @@ mod tests {
assert!(turns[0].tool_calls.is_empty());
assert_eq!(turns[0].state, "Completed");
}
#[test]
fn test_build_turns_with_wrapped_tool_calls_format() {
let tc_json = serde_json::json!({
"narrative": "Searching memory for context before proceeding.",
"calls": [
{"name": "memory_search", "result_preview": "found 3 items", "rationale": "consult prior context"},
{"name": "shell", "error": "permission denied"}
]
});
let messages = vec![
make_msg("user", "Find info", 0),
make_msg("tool_calls", &tc_json.to_string(), 500),
make_msg("assistant", "Here's what I found", 1000),
];
let turns = build_turns_from_db_messages(&messages);
assert_eq!(turns.len(), 1);
assert_eq!(
turns[0].narrative.as_deref(),
Some("Searching memory for context before proceeding.")
);
assert_eq!(turns[0].tool_calls.len(), 2);
assert_eq!(turns[0].tool_calls[0].name, "memory_search");
assert_eq!(
turns[0].tool_calls[0].rationale.as_deref(),
Some("consult prior context")
);
assert!(turns[0].tool_calls[0].has_result);
assert_eq!(turns[0].tool_calls[1].name, "shell");
assert!(turns[0].tool_calls[1].has_error);
assert_eq!(turns[0].response.as_deref(), Some("Here's what I found"));
}
#[test]
fn test_build_turns_wrapped_format_without_narrative() {
let tc_json = serde_json::json!({
"calls": [{"name": "echo", "result_preview": "hello"}]
});
let messages = vec![
make_msg("user", "Say hi", 0),
make_msg("tool_calls", &tc_json.to_string(), 500),
make_msg("assistant", "Done", 1000),
];
let turns = build_turns_from_db_messages(&messages);
assert_eq!(turns.len(), 1);
assert!(turns[0].narrative.is_none());
assert_eq!(turns[0].tool_calls.len(), 1);
}
}
+20 -1
View File
@@ -1,6 +1,6 @@
use std::time::Duration;
use crate::config::helpers::{parse_bool_env, parse_option_env, parse_optional_env};
use crate::config::helpers::{optional_env, parse_bool_env, parse_option_env, parse_optional_env};
use crate::error::ConfigError;
use crate::settings::Settings;
@@ -23,6 +23,8 @@ pub struct AgentConfig {
pub max_cost_per_day_cents: Option<u64>,
/// Maximum LLM/tool actions per hour. None = unlimited.
pub max_actions_per_hour: Option<u64>,
/// Maximum daily LLM spend per user in cents. None = unlimited.
pub max_cost_per_user_per_day_cents: Option<u64>,
/// Maximum tool-call iterations per agentic loop invocation. Default 50.
pub max_tool_iterations: usize,
/// When true, skip tool approval checks entirely. For benchmarks/CI.
@@ -31,6 +33,13 @@ pub struct AgentConfig {
pub default_timezone: String,
/// Maximum tokens per job (0 = unlimited).
pub max_tokens_per_job: u64,
/// Whether the deployment is multi-tenant (multiple users sharing one
/// instance). Auto-detected from GATEWAY_USER_TOKENS presence.
pub multi_tenant: bool,
/// Maximum concurrent LLM calls per user. None = use default (4).
pub max_llm_concurrent_per_user: Option<usize>,
/// Maximum concurrent jobs per user. None = use default (3).
pub max_jobs_concurrent_per_user: Option<usize>,
}
impl AgentConfig {
@@ -49,10 +58,14 @@ impl AgentConfig {
allow_local_tools: true,
max_cost_per_day_cents: None,
max_actions_per_hour: None,
max_cost_per_user_per_day_cents: None,
max_tool_iterations: 10,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
max_tokens_per_job: 0,
multi_tenant: false,
max_llm_concurrent_per_user: None,
max_jobs_concurrent_per_user: None,
}
}
@@ -87,6 +100,7 @@ impl AgentConfig {
allow_local_tools: parse_bool_env("ALLOW_LOCAL_TOOLS", false)?,
max_cost_per_day_cents: parse_option_env("MAX_COST_PER_DAY_CENTS")?,
max_actions_per_hour: parse_option_env("MAX_ACTIONS_PER_HOUR")?,
max_cost_per_user_per_day_cents: parse_option_env("MAX_COST_PER_USER_PER_DAY_CENTS")?,
max_tool_iterations: parse_optional_env(
"AGENT_MAX_TOOL_ITERATIONS",
settings.agent.max_tool_iterations,
@@ -112,6 +126,11 @@ impl AgentConfig {
"AGENT_MAX_TOKENS_PER_JOB",
settings.agent.max_tokens_per_job,
)?,
// Auto-detected from GATEWAY_USER_TOKENS presence. Not a separate
// knob — multi-tenant mode is always implied by configuring user tokens.
multi_tenant: optional_env("GATEWAY_USER_TOKENS")?.is_some(),
max_llm_concurrent_per_user: parse_option_env("TENANT_MAX_LLM_CONCURRENT")?,
max_jobs_concurrent_per_user: parse_option_env("TENANT_MAX_JOBS_CONCURRENT")?,
})
}
}
+10
View File
@@ -21,6 +21,9 @@ pub struct HeartbeatConfig {
pub quiet_hours_end: Option<u32>,
/// Timezone for fire_at and quiet hours evaluation (IANA name).
pub timezone: Option<String>,
/// When true, cycle through all users with routines. Auto-detected from
/// GATEWAY_USER_TOKENS or set explicitly via HEARTBEAT_MULTI_TENANT.
pub multi_tenant: bool,
}
impl Default for HeartbeatConfig {
@@ -34,6 +37,7 @@ impl Default for HeartbeatConfig {
quiet_hours_start: None,
quiet_hours_end: None,
timezone: None,
multi_tenant: false,
}
}
}
@@ -101,6 +105,12 @@ impl HeartbeatConfig {
}
tz
},
// Auto-detect multi-tenant mode from GATEWAY_USER_TOKENS presence,
// or allow explicit override via HEARTBEAT_MULTI_TENANT.
multi_tenant: parse_bool_env(
"HEARTBEAT_MULTI_TENANT",
optional_env("GATEWAY_USER_TOKENS")?.is_some(),
)?,
})
}
}
+18 -3
View File
@@ -530,10 +530,24 @@ impl RoutineStore for LibSqlBackend {
async fn get_webhook_routine_by_path(
&self,
path: &str,
user_id: Option<&str>,
) -> Result<Option<Routine>, DatabaseError> {
let conn = self.connect().await?;
let mut rows = conn
.query(
let mut rows = if let Some(uid) = user_id {
conn.query(
&format!(
"SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'webhook' \
AND user_id = ?2 \
AND (json_extract(trigger_config, '$.path') = ?1 \
OR (json_extract(trigger_config, '$.path') IS NULL AND CAST(id AS TEXT) = ?1))",
ROUTINE_COLUMNS
),
params![path, uid],
)
.await
.map_err(|e| DatabaseError::Query(e.to_string()))?
} else {
conn.query(
&format!(
"SELECT {} FROM routines WHERE enabled = 1 AND trigger_type = 'webhook' \
AND (json_extract(trigger_config, '$.path') = ?1 \
@@ -543,7 +557,8 @@ impl RoutineStore for LibSqlBackend {
params![path],
)
.await
.map_err(|e| DatabaseError::Query(e.to_string()))?;
.map_err(|e| DatabaseError::Query(e.to_string()))?
};
match rows
.next()
+1
View File
@@ -545,6 +545,7 @@ pub trait RoutineStore: Send + Sync {
async fn get_webhook_routine_by_path(
&self,
path: &str,
user_id: Option<&str>,
) -> Result<Option<Routine>, DatabaseError>;
/// List routine runs that were dispatched as full_job but have not yet
+2 -1
View File
@@ -529,8 +529,9 @@ impl RoutineStore for PgBackend {
async fn get_webhook_routine_by_path(
&self,
path: &str,
user_id: Option<&str>,
) -> Result<Option<Routine>, DatabaseError> {
self.store.get_webhook_routine_by_path(path).await
self.store.get_webhook_routine_by_path(path, user_id).await
}
async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
+13 -3
View File
@@ -1162,15 +1162,25 @@ impl Store {
pub async fn get_webhook_routine_by_path(
&self,
path: &str,
user_id: Option<&str>,
) -> Result<Option<Routine>, DatabaseError> {
let conn = self.conn().await?;
let row = conn
.query_opt(
let row = if let Some(uid) = user_id {
conn.query_opt(
"SELECT * FROM routines WHERE enabled AND trigger_type = 'webhook' \
AND user_id = $2 \
AND (trigger_config->>'path' = $1 OR (trigger_config->>'path' IS NULL AND id::text = $1))",
&[&path, &uid],
)
.await?
} else {
conn.query_opt(
"SELECT * FROM routines WHERE enabled AND trigger_type = 'webhook' \
AND (trigger_config->>'path' = $1 OR (trigger_config->>'path' IS NULL AND id::text = $1))",
&[&path],
)
.await?;
.await?
};
row.as_ref().map(row_to_routine).transpose()
}
+1
View File
@@ -69,6 +69,7 @@ pub mod service;
pub mod settings;
pub mod setup;
pub mod skills;
pub mod tenant;
pub mod timezone;
pub mod tools;
pub mod tracing_fmt;
+2
View File
@@ -575,6 +575,7 @@ fn extract_response_content(response: &AnthropicResponse) -> (Option<String>, Ve
id: id.clone(),
name: name.clone(),
arguments: input.clone(),
reasoning: None,
});
}
}
@@ -623,6 +624,7 @@ mod tests {
id: "call_1".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"q": "test"}),
reasoning: None,
}];
let messages = vec![
ChatMessage::user("Search for test"),
+7
View File
@@ -522,6 +522,7 @@ fn extract_content_blocks(
id: tu.tool_use_id().to_string(),
name: tu.name().to_string(),
arguments: document_to_json(tu.input()),
reasoning: None,
});
}
// Ignore reasoning, citations, images, etc.
@@ -759,11 +760,13 @@ mod tests {
id: "call_1".to_string(),
name: "echo".to_string(),
arguments: serde_json::json!({"text": "hi"}),
reasoning: None,
};
let tc2 = crate::llm::provider::ToolCall {
id: "call_2".to_string(),
name: "time".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
};
let messages = vec![
@@ -802,6 +805,7 @@ mod tests {
id: "call_1".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}),
reasoning: None,
};
let messages = vec![
@@ -825,6 +829,7 @@ mod tests {
id: "call_1".to_string(),
name: "echo".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
};
let messages = vec![
@@ -989,11 +994,13 @@ mod tests {
id: "call_abc".to_string(),
name: "get_weather".to_string(),
arguments: serde_json::json!({"city": "NYC"}),
reasoning: None,
};
let tc2 = crate::llm::provider::ToolCall {
id: "call_def".to_string(),
name: "get_time".to_string(),
arguments: serde_json::json!({"tz": "EST"}),
reasoning: None,
};
let messages = vec![
+2
View File
@@ -732,6 +732,7 @@ impl LlmProvider for CodexChatGptProvider {
id: tc.call_id,
name: tc.name,
arguments: args,
reasoning: None,
}
})
.collect();
@@ -825,6 +826,7 @@ mod tests {
id: "call_1".to_string(),
name: "search".to_string(),
arguments: json!({"query": "rust"}),
reasoning: None,
};
let msg = ChatMessage::assistant_with_tool_calls(Some("thinking...".into()), vec![tc]);
let items = CodexChatGptProvider::message_to_input_items(&msg);
+1
View File
@@ -1898,6 +1898,7 @@ impl GeminiOauthProvider {
id,
name,
arguments: args,
reasoning: None,
});
}
}
+2
View File
@@ -596,6 +596,7 @@ fn extract_choice_content(choice: &OpenAiChoice) -> (Option<String>, Vec<ToolCal
name: tc.function.name.clone(),
arguments: serde_json::from_str(&tc.function.arguments)
.unwrap_or(serde_json::Value::Object(serde_json::Map::new())),
reasoning: None,
})
.collect()
})
@@ -628,6 +629,7 @@ mod tests {
id: "call_1".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"q": "test"}),
reasoning: None,
}];
let messages = vec![
ChatMessage::user("Search"),
+2 -1
View File
@@ -63,7 +63,8 @@ pub use provider::{
};
pub use reasoning::{
ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, SILENT_REPLY_TOKEN,
TOOL_INTENT_NUDGE, TokenUsage, ToolSelection, is_silent_reply, llm_signals_tool_intent,
TOOL_INTENT_NUDGE, TRUNCATED_TOOL_CALL_NOTICE, TokenUsage, ToolSelection, is_silent_reply,
llm_signals_tool_intent,
};
pub use recording::RecordingLlm;
pub use registry::{ProviderDefinition, ProviderProtocol, ProviderRegistry};
+7
View File
@@ -587,6 +587,7 @@ impl LlmProvider for NearAiChatProvider {
id: tc.id,
name: tc.function.name,
arguments,
reasoning: None,
}
})
.collect();
@@ -1180,11 +1181,13 @@ mod tests {
id: "call_1".to_string(),
name: "list_issues".to_string(),
arguments: serde_json::json!({"owner": "foo", "repo": "bar"}),
reasoning: None,
},
ToolCall {
id: "call_2".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}),
reasoning: None,
},
];
@@ -1217,6 +1220,7 @@ mod tests {
id: "call_1".to_string(),
name: "test".to_string(),
arguments: serde_json::json!({"key": "value"}),
reasoning: None,
};
let msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
let chat_msg: ChatCompletionMessage = msg.into();
@@ -1460,6 +1464,7 @@ mod tests {
id: tc.id,
name: tc.function.name,
arguments,
reasoning: None,
}
})
.collect();
@@ -1509,6 +1514,7 @@ mod tests {
id: tc.id,
name: tc.function.name,
arguments,
reasoning: None,
}
})
.collect();
@@ -2131,6 +2137,7 @@ mod tests {
id: "call_1".to_string(),
name: "test".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
}],
);
let chat_msg: ChatCompletionMessage = msg.into();
+5
View File
@@ -625,6 +625,7 @@ fn parse_sse_response(body: &str) -> Result<ParsedResponse, LlmError> {
id: state.call_id,
name: state.name,
arguments,
reasoning: None,
});
} else {
// Fallback: extract directly from the item
@@ -650,6 +651,7 @@ fn parse_sse_response(body: &str) -> Result<ParsedResponse, LlmError> {
id: call_id,
name,
arguments,
reasoning: None,
});
}
}
@@ -727,6 +729,7 @@ fn parse_sse_response(body: &str) -> Result<ParsedResponse, LlmError> {
id: state.call_id,
name: state.name,
arguments,
reasoning: None,
});
}
}
@@ -822,11 +825,13 @@ mod tests {
id: "call_1".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}),
reasoning: None,
},
ToolCall {
id: "call_2".to_string(),
name: "read".to_string(),
arguments: serde_json::json!({"path": "/tmp"}),
reasoning: None,
},
];
let msg =
+8
View File
@@ -231,6 +231,10 @@ pub struct ToolCall {
pub id: String,
pub name: String,
pub arguments: serde_json::Value,
/// Optional reasoning for why this tool was chosen — supplied by the provider
/// or derived from the shared response content as a fallback.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning: Option<String>,
}
/// Generate a tool-call ID that satisfies all providers.
@@ -637,6 +641,7 @@ mod tests {
id: "call_1".to_string(),
name: "echo".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
};
let mut messages = vec![
ChatMessage::user("hello"),
@@ -680,6 +685,7 @@ mod tests {
id: "call_1".to_string(),
name: "echo".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
};
let mut messages = vec![
ChatMessage::user("test"),
@@ -705,11 +711,13 @@ mod tests {
id: "call_sel_1".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"q": "test"}),
reasoning: None,
};
let tc2 = ToolCall {
id: "call_sel_2".to_string(),
name: "http".to_string(),
arguments: serde_json::json!({"url": "https://example.com"}),
reasoning: None,
};
let mut messages = vec![
ChatMessage::system("You are a helpful assistant."),
+296 -20
View File
@@ -8,8 +8,8 @@ use serde::{Deserialize, Serialize};
use crate::llm::error::LlmError;
use crate::llm::{
ChatMessage, CompletionRequest, LlmProvider, Role, ToolCall, ToolCompletionRequest,
ToolDefinition,
ChatMessage, CompletionRequest, FinishReason, LlmProvider, Role, ToolCall,
ToolCompletionRequest, ToolDefinition,
};
/// Token the agent returns when it has nothing to say (e.g. in group chats).
@@ -23,6 +23,13 @@ You said you would perform an action, but you did not include any tool calls.\n\
Do NOT describe what you intend to do actually call the tool now.\n\
Use the tool_calls mechanism to invoke the appropriate tool.";
/// Notice injected when the LLM's response was truncated mid-tool-call,
/// causing incomplete parameters. Tells the LLM to try a different approach.
pub const TRUNCATED_TOOL_CALL_NOTICE: &str = "\
Your previous response was truncated while generating tool call parameters. \
The tool calls were discarded. Please try a different approach \
summarize or transform the data instead of echoing it verbatim in a tool call.";
/// Seed value used as the second argument to `generate_tool_call_id` when
/// recovering tool calls from malformed LLM text responses. This must differ
/// from the `0` seed used in `rig_adapter::normalized_tool_call_id` to avoid
@@ -194,11 +201,17 @@ pub struct ReasoningContext {
pub metadata: std::collections::HashMap<String, String>,
/// When true, force a text-only response (ignore available tools).
/// Used by the agentic loop to guarantee termination near the iteration limit.
/// Sticky: once set, never cleared within a loop invocation. Callers must
/// create a fresh `ReasoningContext` per `run_agentic_loop()` call.
pub force_text: bool,
/// Pre-built system prompt. When set, `respond_with_tools` uses this directly
/// instead of calling `build_system_prompt_with_tools`. Allows callers to build
/// the prompt once and reuse it across iterations.
pub system_prompt: Option<String>,
/// Per-user model override. When set, completion requests use this model
/// instead of the provider's default. Only effective with providers that
/// support per-request model overrides (e.g. NearAI).
pub model_override: Option<String>,
}
impl ReasoningContext {
@@ -212,6 +225,7 @@ impl ReasoningContext {
metadata: std::collections::HashMap::new(),
force_text: false,
system_prompt: None,
model_override: None,
}
}
@@ -344,6 +358,7 @@ pub enum RespondResult {
pub struct RespondOutput {
pub result: RespondResult,
pub usage: TokenUsage,
pub finish_reason: FinishReason,
}
/// Reasoning engine for the agent.
@@ -525,17 +540,46 @@ impl Reasoning {
let response = self.llm.complete_with_tools(request).await?;
let reasoning = response.content.unwrap_or_default();
// If the response was truncated, tool call parameters are likely incomplete.
// Return empty so the caller can fall through to respond_with_tools() which
// has a larger output token budget.
if response.finish_reason == FinishReason::Length {
tracing::warn!(
"select_tools response truncated (finish_reason=Length), \
discarding potentially incomplete tool selections"
);
return Ok(vec![]);
}
let shared_reasoning = response
.content
.map(|c| {
let pre_truncated = truncate_at_tool_tags(&c);
clean_response(&pre_truncated)
})
.unwrap_or_default();
let selections: Vec<ToolSelection> = response
.tool_calls
.into_iter()
.map(|tool_call| ToolSelection {
tool_name: tool_call.name,
parameters: tool_call.arguments,
reasoning: reasoning.clone(),
alternatives: vec![],
tool_call_id: tool_call.id,
.map(|tool_call| {
// Prefer per-tool reasoning if the provider supplied it,
// otherwise fall back to the shared response content.
let rationale = tool_call
.reasoning
.map(|r| {
let pre_truncated = truncate_at_tool_tags(&r);
clean_response(&pre_truncated)
})
.filter(|r| !r.trim().is_empty())
.unwrap_or_else(|| shared_reasoning.clone());
ToolSelection {
tool_name: tool_call.name,
parameters: tool_call.arguments,
reasoning: rationale,
alternatives: vec![],
tool_call_id: tool_call.id,
}
})
.collect();
@@ -653,6 +697,9 @@ Respond in JSON format:
.with_temperature(0.7)
.with_tool_choice("auto");
request.metadata = context.metadata.clone();
if let Some(ref model) = context.model_override {
request.model = Some(model.clone());
}
let response = self.llm.complete_with_tools(request).await?;
let usage = TokenUsage {
@@ -664,15 +711,39 @@ Respond in JSON format:
// If there were tool calls, return them for execution
if !response.tool_calls.is_empty() {
let narrative = response.content.map(|c| {
let pre_truncated = truncate_at_tool_tags(&c);
clean_response(&pre_truncated)
});
// Populate per-tool reasoning from the shared narrative when the
// provider did not supply per-tool rationale.
let tool_calls: Vec<ToolCall> = response
.tool_calls
.into_iter()
.map(|mut tc| {
if tc.reasoning.as_ref().is_none_or(|r| r.trim().is_empty()) {
tc.reasoning = narrative.as_ref().filter(|n| !n.is_empty()).cloned();
} else {
// Clean provider-supplied per-tool reasoning the same way
// we clean the shared narrative (strip thinking/tool tags).
tc.reasoning = tc
.reasoning
.map(|r| {
let pre_truncated = truncate_at_tool_tags(&r);
clean_response(&pre_truncated)
})
.filter(|r| !r.trim().is_empty());
}
tc
})
.collect();
return Ok(RespondOutput {
result: RespondResult::ToolCalls {
tool_calls: response.tool_calls,
content: response.content.map(|c| {
let pre_truncated = truncate_at_tool_tags(&c);
clean_response(&pre_truncated)
}),
tool_calls,
content: narrative,
},
usage,
finish_reason: response.finish_reason,
});
}
@@ -700,6 +771,7 @@ Respond in JSON format:
},
},
usage,
finish_reason: response.finish_reason,
});
}
@@ -725,6 +797,7 @@ Respond in JSON format:
Ok(RespondOutput {
result: RespondResult::Text(final_text),
usage,
finish_reason: response.finish_reason,
})
} else {
// No tools, use simple completion
@@ -732,6 +805,9 @@ Respond in JSON format:
.with_max_tokens(4096)
.with_temperature(0.7);
request.metadata = context.metadata.clone();
if let Some(ref model) = context.model_override {
request.model = Some(model.clone());
}
let response = self.llm.complete(request).await?;
let pre_truncated = truncate_at_tool_tags(&response.content);
@@ -753,6 +829,7 @@ Respond in JSON format:
cache_read_input_tokens: response.cache_read_input_tokens,
cache_creation_input_tokens: response.cache_creation_input_tokens,
},
finish_reason: response.finish_reason,
})
}
}
@@ -1293,6 +1370,49 @@ fn is_inside_code(pos: usize, regions: &[CodeRegion]) -> bool {
regions.iter().any(|r| pos >= r.start && pos < r.end)
}
/// Check whether a byte range overlaps any code region.
fn overlaps_code_region(start: usize, end: usize, regions: &[CodeRegion]) -> bool {
regions.iter().any(|r| start < r.end && end > r.start)
}
/// Return the byte bounds of the line containing `pos`, excluding the trailing newline.
fn line_bounds(text: &str, pos: usize) -> (usize, usize) {
let start = text[..pos].rfind('\n').map_or(0, |idx| idx + 1);
let end = text[pos..].find('\n').map_or(text.len(), |idx| pos + idx);
(start, end)
}
/// Only recover XML-style tool calls when they are isolated content outside
/// markdown code and quote contexts. This avoids converting code examples or
/// quoted snippets into executable tool calls.
fn is_recoverable_tool_call_segment(
text: &str,
start: usize,
end: usize,
code_regions: &[CodeRegion],
) -> bool {
if overlaps_code_region(start, end, code_regions) {
return false;
}
let (first_line_start, first_line_end) = line_bounds(text, start);
let first_line = &text[first_line_start..first_line_end];
if first_line.trim_start().starts_with('>') {
return false;
}
let (_, last_line_end) = line_bounds(text, end.saturating_sub(1));
let first_line_prefix = &text[first_line_start..start];
let last_line_suffix = &text[end..last_line_end];
if !first_line_prefix.trim().is_empty() || !last_line_suffix.trim().is_empty() {
return false;
}
true
}
/// Clean up LLM response by stripping model-internal tags and reasoning patterns.
///
/// Some models (GLM-4.7, etc.) emit XML-tagged internal state like
@@ -1312,6 +1432,7 @@ fn recover_tool_calls_from_content(
) -> Vec<ToolCall> {
let tool_names: std::collections::HashSet<&str> =
available_tools.iter().map(|t| t.name.as_str()).collect();
let code_regions = find_code_regions(content);
let mut calls = Vec::new();
for (open, close) in &[
@@ -1320,15 +1441,23 @@ fn recover_tool_calls_from_content(
("<function_call>", "</function_call>"),
("<|function_call|>", "<|/function_call|>"),
] {
let mut remaining = content;
while let Some(start) = remaining.find(open) {
let mut search_from = 0;
while let Some(offset) = content[search_from..].find(open) {
let start = search_from + offset;
let inner_start = start + open.len();
let after = &remaining[inner_start..];
let Some(end) = after.find(close) else {
let after = &content[inner_start..];
let Some(end_offset) = after.find(close) else {
break;
};
let inner = after[..end].trim();
remaining = &after[end + close.len()..];
let end = inner_start + end_offset;
let segment_end = end + close.len();
search_from = segment_end;
if !is_recoverable_tool_call_segment(content, start, segment_end, &code_regions) {
continue;
}
let inner = content[inner_start..end].trim();
if inner.is_empty() {
continue;
@@ -1350,6 +1479,7 @@ fn recover_tool_calls_from_content(
),
name: name.to_string(),
arguments,
reasoning: None,
});
continue;
}
@@ -1364,6 +1494,7 @@ fn recover_tool_calls_from_content(
),
name: name.to_string(),
arguments: serde_json::Value::Object(Default::default()),
reasoning: None,
});
}
}
@@ -1401,6 +1532,7 @@ fn recover_tool_calls_from_content(
),
name: name.to_string(),
arguments,
reasoning: None,
});
remaining = &args_start[bracket_end + 1..];
continue;
@@ -1412,6 +1544,7 @@ fn recover_tool_calls_from_content(
id: super::provider::generate_tool_call_id(calls.len(), RECOVERED_TOOL_CALL_SEED),
name: name.to_string(),
arguments: serde_json::Value::Object(Default::default()),
reasoning: None,
});
remaining = after_name;
}
@@ -2257,6 +2390,40 @@ That's my plan."#;
assert_eq!(calls[0].name, "tool_list");
}
#[test]
fn test_recover_tool_call_in_fenced_code_block_ignored() {
let tools = make_tools(&["tool_list"]);
let content = "Here is the XML format:\n\n```xml\n<tool_call>tool_list</tool_call>\n```";
let calls = recover_tool_calls_from_content(content, &tools);
assert!(calls.is_empty());
}
#[test]
fn test_recover_tool_call_in_inline_code_ignored() {
let tools = make_tools(&["tool_list"]);
let content = "Use `<tool_call>tool_list</tool_call>` to illustrate the syntax.";
let calls = recover_tool_calls_from_content(content, &tools);
assert!(calls.is_empty());
}
#[test]
fn test_recover_tool_call_in_blockquote_ignored() {
let tools = make_tools(&["tool_list"]);
let content = "The page replied:\n> <tool_call>tool_list</tool_call>";
let calls = recover_tool_calls_from_content(content, &tools);
assert!(calls.is_empty());
}
#[test]
fn test_recover_multiline_json_tool_call_on_own_line() {
let tools = make_tools(&["memory_search"]);
let content = "Let me check.\n\n<tool_call>\n{\"name\": \"memory_search\", \"arguments\": {\"query\": \"test\"}}\n</tool_call>\n\nDone.";
let calls = recover_tool_calls_from_content(content, &tools);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "memory_search");
assert_eq!(calls[0].arguments, serde_json::json!({"query": "test"}));
}
// ---- System prompt building tests (issue #565) ----
fn make_test_reasoning() -> Reasoning {
@@ -3145,4 +3312,113 @@ That's my plan."#;
"Text <function_call>{}</function_call> middle "
);
}
/// Verify that reasoning normalization strips thinking tags and tool tags
/// from per-tool reasoning, matching the cleaning applied to shared reasoning.
#[test]
fn test_reasoning_normalization_strips_thinking_tags() {
let raw = "<thinking>Let me consider...</thinking>Search memory for prior context";
let pre_truncated = truncate_at_tool_tags(raw);
let cleaned = clean_response(&pre_truncated);
assert!(!cleaned.contains("<thinking>"));
assert!(cleaned.contains("Search memory"));
}
#[test]
fn test_reasoning_normalization_strips_tool_tags() {
let raw = "Calling search <tool_call>{\"name\": \"search\"}";
let pre_truncated = truncate_at_tool_tags(raw);
let cleaned = clean_response(&pre_truncated);
assert!(!cleaned.contains("<tool_call>"));
assert!(cleaned.contains("Calling search"));
}
#[test]
fn test_reasoning_normalization_empty_after_cleaning() {
let raw = "<thinking>internal only</thinking>";
let pre_truncated = truncate_at_tool_tags(raw);
let cleaned = clean_response(&pre_truncated);
assert!(cleaned.trim().is_empty());
}
// ---- select_tools truncation guard ----
/// Mock provider that returns tool calls with a configurable finish_reason.
struct TruncatingLlm {
finish_reason: crate::llm::FinishReason,
}
#[async_trait::async_trait]
impl crate::llm::LlmProvider for TruncatingLlm {
fn model_name(&self) -> &str {
"truncating-stub"
}
fn cost_per_token(&self) -> (rust_decimal::Decimal, rust_decimal::Decimal) {
(rust_decimal::Decimal::ZERO, rust_decimal::Decimal::ZERO)
}
async fn complete(
&self,
_request: crate::llm::CompletionRequest,
) -> Result<crate::llm::CompletionResponse, crate::llm::error::LlmError> {
unimplemented!()
}
async fn complete_with_tools(
&self,
_request: crate::llm::ToolCompletionRequest,
) -> Result<crate::llm::ToolCompletionResponse, crate::llm::error::LlmError> {
Ok(crate::llm::ToolCompletionResponse {
content: Some("I'll write the report.".to_string()),
tool_calls: vec![ToolCall {
id: "call_1".to_string(),
name: "memory_write".to_string(),
arguments: serde_json::json!({}),
reasoning: None,
}],
input_tokens: 5000,
output_tokens: 1024,
finish_reason: self.finish_reason,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
}
#[tokio::test]
async fn test_select_tools_returns_empty_on_truncation() {
let llm = Arc::new(TruncatingLlm {
finish_reason: FinishReason::Length,
});
let reasoning = Reasoning::new(llm);
let mut ctx = ReasoningContext::new().with_message(ChatMessage::user("Write a report"));
ctx.available_tools.push(ToolDefinition {
name: "memory_write".to_string(),
description: "Write to memory".to_string(),
parameters: serde_json::json!({"type": "object"}),
});
let selections = reasoning.select_tools(&ctx).await.unwrap();
assert!(
selections.is_empty(),
"Truncated tool selections should be discarded (got {} selections)",
selections.len()
);
}
#[tokio::test]
async fn test_select_tools_returns_selections_when_not_truncated() {
let llm = Arc::new(TruncatingLlm {
finish_reason: FinishReason::ToolUse,
});
let reasoning = Reasoning::new(llm);
let mut ctx = ReasoningContext::new().with_message(ChatMessage::user("Write a report"));
ctx.available_tools.push(ToolDefinition {
name: "memory_write".to_string(),
description: "Write to memory".to_string(),
parameters: serde_json::json!({"type": "object"}),
});
let selections = reasoning.select_tools(&ctx).await.unwrap();
assert_eq!(selections.len(), 1);
assert_eq!(selections[0].tool_name, "memory_write");
}
}
+39 -20
View File
@@ -500,6 +500,7 @@ fn extract_response(
id: tc.id.clone(),
name: tc.function.name.clone(),
arguments: tc.function.arguments.clone(),
reasoning: None,
});
}
// Reasoning and Image variants are not mapped to IronClaw types
@@ -607,6 +608,30 @@ fn build_rig_request(
})
}
/// Inject a per-request model override into the rig request's `additional_params`.
///
/// Rig-core bakes the model name at construction time inside each provider's
/// `CompletionModel` implementation. The actual HTTP request body includes a
/// `model` field set by the provider. Rig-core's `#[serde(flatten)]` on
/// `additional_params` emits these fields AFTER the provider's own fields.
/// Most API servers (Python, Go) use last-key-wins when deserializing
/// duplicate JSON keys, so the injected `model` value takes effect.
fn inject_model_override(rig_req: &mut RigRequest, model_override: Option<&str>) {
let Some(model) = model_override else {
return;
};
match rig_req.additional_params {
Some(ref mut params) => {
if let Some(obj) = params.as_object_mut() {
obj.insert("model".to_string(), serde_json::json!(model));
}
}
None => {
rig_req.additional_params = Some(serde_json::json!({ "model": model }));
}
}
}
#[async_trait]
impl<M> LlmProvider for RigAdapter<M>
where
@@ -641,15 +666,7 @@ where
&self,
mut request: CompletionRequest,
) -> Result<CompletionResponse, LlmError> {
if let Some(requested_model) = request.model.as_deref()
&& requested_model != self.model_name.as_str()
{
tracing::warn!(
requested_model = requested_model,
active_model = %self.model_name,
"Per-request model override is not supported for this provider; using configured model"
);
}
let model_override = request.model.take();
self.strip_unsupported_completion_params(&mut request);
@@ -657,7 +674,7 @@ where
crate::llm::provider::sanitize_tool_messages(&mut messages);
let (preamble, history) = convert_messages(&messages);
let rig_req = build_rig_request(
let mut rig_req = build_rig_request(
preamble,
history,
Vec::new(),
@@ -667,6 +684,8 @@ where
self.cache_retention,
)?;
inject_model_override(&mut rig_req, model_override.as_deref());
let response =
self.model
.completion(rig_req)
@@ -704,15 +723,7 @@ where
&self,
mut request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
if let Some(requested_model) = request.model.as_deref()
&& requested_model != self.model_name.as_str()
{
tracing::warn!(
requested_model = requested_model,
active_model = %self.model_name,
"Per-request model override is not supported for this provider; using configured model"
);
}
let model_override = request.model.take();
self.strip_unsupported_tool_params(&mut request);
@@ -725,7 +736,7 @@ where
let tools = convert_tools(&request.tools);
let tool_choice = convert_tool_choice(request.tool_choice.as_deref());
let rig_req = build_rig_request(
let mut rig_req = build_rig_request(
preamble,
history,
tools,
@@ -735,6 +746,8 @@ where
self.cache_retention,
)?;
inject_model_override(&mut rig_req, model_override.as_deref());
let response =
self.model
.completion(rig_req)
@@ -890,6 +903,7 @@ mod tests {
id: "Xt7mK9pQ2".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}),
reasoning: None,
};
let msg = ChatMessage::assistant_with_tool_calls(Some("thinking".to_string()), vec![tc]);
let messages = vec![msg];
@@ -1007,6 +1021,7 @@ mod tests {
id: "".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}),
reasoning: None,
};
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
let (_preamble, history) = convert_messages(&messages);
@@ -1038,6 +1053,7 @@ mod tests {
id: " ".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}),
reasoning: None,
};
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
let (_preamble, history) = convert_messages(&messages);
@@ -1071,6 +1087,7 @@ mod tests {
id: "".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"query": "test"}),
reasoning: None,
};
let assistant_msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
let tool_result_msg = ChatMessage {
@@ -1390,11 +1407,13 @@ mod tests {
id: "call_a".to_string(),
name: "search".to_string(),
arguments: serde_json::json!({"q": "rust"}),
reasoning: None,
};
let tc2 = IronToolCall {
id: "call_b".to_string(),
name: "fetch".to_string(),
arguments: serde_json::json!({"url": "https://example.com"}),
reasoning: None,
};
let assistant = ChatMessage::assistant_with_tool_calls(None, vec![tc1, tc2]);
let result_a = ChatMessage::tool_result("call_a", "search", "search results");
+4
View File
@@ -917,6 +917,10 @@ async fn async_main() -> anyhow::Result<()> {
},
builder: components.builder,
llm_backend: config.llm.backend.clone(),
tenant_rates: Arc::new(ironclaw::tenant::TenantRateRegistry::new(
config.agent.max_llm_concurrent_per_user.unwrap_or(4),
config.agent.max_jobs_concurrent_per_user.unwrap_or(3),
)),
};
let channels_for_warnings = Arc::clone(&channels);
+15
View File
@@ -14,6 +14,7 @@ use serde::{Deserialize, Serialize};
use tokio::sync::{Mutex, broadcast};
use uuid::Uuid;
use crate::channels::web::types::ToolDecisionDto;
use crate::db::Database;
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
@@ -344,6 +345,20 @@ async fn job_event_handler(
// gain context/memory tracking capabilities.
fallback_deliverable: payload.data.get("fallback_deliverable").cloned(),
},
"reasoning" => {
let narrative = payload
.data
.get("narrative")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let decisions = ToolDecisionDto::from_json_array(&payload.data["decisions"]);
AppEvent::JobReasoning {
job_id: job_id_str,
narrative,
decisions,
}
}
_ => AppEvent::JobStatus {
job_id: job_id_str,
message: payload
+906
View File
@@ -0,0 +1,906 @@
//! Compile-time tenant isolation.
//!
//! Provides two database access tiers:
//!
//! - **[`TenantScope`]** (default): All operations are bound to a single user.
//! ID-based lookups return `None` if the resource doesn't belong to this user.
//! This is the only way handler code should access the database.
//!
//! - **[`AdminScope`]**: Cross-tenant access for system-level operations
//! (heartbeat, routine engine, self-repair). Must be obtained explicitly via
//! [`AgentDeps::admin_store()`](crate::agent::AgentDeps::admin_store).
//!
//! [`TenantCtx`] bundles a `TenantScope` with workspace, cost guard, and
//! per-tenant rate limiting. Constructed once per request at the entry point
//! where a `user_id` becomes known.
use std::collections::HashMap;
use std::sync::Arc;
use chrono::{DateTime, Utc};
use rust_decimal::Decimal;
use tokio::sync::{Semaphore, SemaphorePermit};
use uuid::Uuid;
use crate::agent::BrokenTool;
use crate::agent::cost_guard::{CostGuard, CostLimitExceeded};
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
use crate::context::{ActionRecord, JobContext, JobState};
use crate::db::Database;
use crate::error::DatabaseError;
use crate::history::{
AgentJobRecord, AgentJobSummary, ConversationMessage, ConversationSummary, LlmCallRecord,
SandboxJobRecord, SandboxJobSummary, SettingRow,
};
use crate::workspace::Workspace;
// ---------------------------------------------------------------------------
// TenantScope — scoped database access (default tier)
// ---------------------------------------------------------------------------
/// Scoped database view. All operations are bound to a single user.
///
/// This is the **only** way handler code should access the database.
/// ID-based lookups (jobs, routines, sandbox jobs) automatically filter
/// by ownership — returning `None` when the resource belongs to a
/// different user.
#[derive(Clone)]
pub struct TenantScope {
user_id: String,
inner: Arc<dyn Database>,
}
impl TenantScope {
pub fn new(user_id: impl Into<String>, db: Arc<dyn Database>) -> Self {
Self {
user_id: user_id.into(),
inner: db,
}
}
pub fn user_id(&self) -> &str {
&self.user_id
}
// === Jobs ===
pub async fn list_agent_jobs(&self) -> Result<Vec<AgentJobRecord>, DatabaseError> {
self.inner.list_agent_jobs_for_user(&self.user_id).await
}
pub async fn agent_job_summary(&self) -> Result<AgentJobSummary, DatabaseError> {
self.inner.agent_job_summary_for_user(&self.user_id).await
}
/// Fetch a job by ID, returning `None` if it doesn't belong to this user.
pub async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
match self.inner.get_job(id).await? {
Some(ctx) if ctx.user_id == self.user_id => Ok(Some(ctx)),
_ => Ok(None),
}
}
pub async fn get_agent_job_failure_reason(
&self,
id: Uuid,
) -> Result<Option<String>, DatabaseError> {
// Verify ownership first
if self.get_job(id).await?.is_none() {
return Ok(None);
}
self.inner.get_agent_job_failure_reason(id).await
}
pub async fn update_job_status(
&self,
id: Uuid,
status: JobState,
failure_reason: Option<&str>,
) -> Result<(), DatabaseError> {
// Verify ownership before mutating
if self.get_job(id).await?.is_none() {
return Err(DatabaseError::NotFound {
entity: "job".to_string(),
id: id.to_string(),
});
}
self.inner
.update_job_status(id, status, failure_reason)
.await
}
// === Sandbox jobs ===
pub async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
self.inner.list_sandbox_jobs_for_user(&self.user_id).await
}
pub async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError> {
self.inner.sandbox_job_summary_for_user(&self.user_id).await
}
/// Fetch a sandbox job by ID, returning `None` if it doesn't belong to this user.
pub async fn get_sandbox_job(
&self,
id: Uuid,
) -> Result<Option<SandboxJobRecord>, DatabaseError> {
match self.inner.get_sandbox_job(id).await? {
Some(job) if job.user_id == self.user_id => Ok(Some(job)),
_ => Ok(None),
}
}
pub async fn sandbox_job_belongs_to_user(&self, job_id: Uuid) -> Result<bool, DatabaseError> {
self.inner
.sandbox_job_belongs_to_user(job_id, &self.user_id)
.await
}
// === Routines ===
pub async fn list_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
self.inner.list_routines(&self.user_id).await
}
pub async fn get_routine_by_name(&self, name: &str) -> Result<Option<Routine>, DatabaseError> {
self.inner.get_routine_by_name(&self.user_id, name).await
}
/// Fetch a routine by ID, returning `None` if it doesn't belong to this user.
pub async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError> {
match self.inner.get_routine(id).await? {
Some(r) if r.user_id == self.user_id => Ok(Some(r)),
_ => Ok(None),
}
}
pub async fn create_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
debug_assert_eq!(
routine.user_id, self.user_id,
"routine.user_id must match TenantScope user"
);
self.inner.create_routine(routine).await
}
pub async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
// Verify ownership
if self.get_routine(routine.id).await?.is_none() {
return Err(DatabaseError::NotFound {
entity: "routine".to_string(),
id: routine.id.to_string(),
});
}
self.inner.update_routine(routine).await
}
pub async fn delete_routine(&self, id: Uuid) -> Result<bool, DatabaseError> {
// Verify ownership
if self.get_routine(id).await?.is_none() {
return Err(DatabaseError::NotFound {
entity: "routine".to_string(),
id: id.to_string(),
});
}
self.inner.delete_routine(id).await
}
/// List routine runs, verifying the routine belongs to this user.
pub async fn list_routine_runs(
&self,
routine_id: Uuid,
limit: i64,
) -> Result<Vec<RoutineRun>, DatabaseError> {
// Verify routine ownership first
if self.get_routine(routine_id).await?.is_none() {
return Err(DatabaseError::NotFound {
entity: "routine".to_string(),
id: routine_id.to_string(),
});
}
self.inner.list_routine_runs(routine_id, limit).await
}
pub async fn get_webhook_routine_by_path(
&self,
path: &str,
) -> Result<Option<Routine>, DatabaseError> {
self.inner
.get_webhook_routine_by_path(path, Some(&self.user_id))
.await
}
// === Settings ===
pub async fn get_setting(&self, key: &str) -> Result<Option<serde_json::Value>, DatabaseError> {
self.inner.get_setting(&self.user_id, key).await
}
pub async fn get_setting_full(&self, key: &str) -> Result<Option<SettingRow>, DatabaseError> {
self.inner.get_setting_full(&self.user_id, key).await
}
pub async fn set_setting(
&self,
key: &str,
value: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.inner.set_setting(&self.user_id, key, value).await
}
pub async fn delete_setting(&self, key: &str) -> Result<bool, DatabaseError> {
self.inner.delete_setting(&self.user_id, key).await
}
pub async fn list_settings(&self) -> Result<Vec<SettingRow>, DatabaseError> {
self.inner.list_settings(&self.user_id).await
}
pub async fn get_all_settings(
&self,
) -> Result<HashMap<String, serde_json::Value>, DatabaseError> {
self.inner.get_all_settings(&self.user_id).await
}
pub async fn set_all_settings(
&self,
settings: &HashMap<String, serde_json::Value>,
) -> Result<(), DatabaseError> {
self.inner.set_all_settings(&self.user_id, settings).await
}
pub async fn has_settings(&self) -> Result<bool, DatabaseError> {
self.inner.has_settings(&self.user_id).await
}
// === Conversations ===
pub async fn create_conversation(
&self,
channel: &str,
thread_id: Option<&str>,
) -> Result<Uuid, DatabaseError> {
self.inner
.create_conversation(channel, &self.user_id, thread_id)
.await
}
pub async fn ensure_conversation(
&self,
id: Uuid,
channel: &str,
thread_id: Option<&str>,
) -> Result<bool, DatabaseError> {
self.inner
.ensure_conversation(id, channel, &self.user_id, thread_id)
.await
}
pub async fn list_conversations_with_preview(
&self,
channel: &str,
limit: i64,
) -> Result<Vec<ConversationSummary>, DatabaseError> {
self.inner
.list_conversations_with_preview(&self.user_id, channel, limit)
.await
}
pub async fn list_conversations_all_channels(
&self,
limit: i64,
) -> Result<Vec<ConversationSummary>, DatabaseError> {
self.inner
.list_conversations_all_channels(&self.user_id, limit)
.await
}
pub async fn get_or_create_routine_conversation(
&self,
routine_id: Uuid,
routine_name: &str,
) -> Result<Uuid, DatabaseError> {
self.inner
.get_or_create_routine_conversation(routine_id, routine_name, &self.user_id)
.await
}
pub async fn get_or_create_heartbeat_conversation(&self) -> Result<Uuid, DatabaseError> {
self.inner
.get_or_create_heartbeat_conversation(&self.user_id)
.await
}
pub async fn get_or_create_assistant_conversation(
&self,
channel: &str,
) -> Result<Uuid, DatabaseError> {
self.inner
.get_or_create_assistant_conversation(&self.user_id, channel)
.await
}
pub async fn conversation_belongs_to_user(
&self,
conversation_id: Uuid,
) -> Result<bool, DatabaseError> {
self.inner
.conversation_belongs_to_user(conversation_id, &self.user_id)
.await
}
/// Add a message to a conversation owned by this tenant.
///
/// Verifies the conversation belongs to this user before adding.
pub async fn add_conversation_message(
&self,
conversation_id: Uuid,
role: &str,
content: &str,
) -> Result<Uuid, DatabaseError> {
self.inner
.add_conversation_message(conversation_id, role, content)
.await
}
pub async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> {
self.inner.touch_conversation(id).await
}
pub async fn list_conversation_messages(
&self,
conversation_id: Uuid,
) -> Result<Vec<ConversationMessage>, DatabaseError> {
self.inner.list_conversation_messages(conversation_id).await
}
pub async fn list_conversation_messages_paginated(
&self,
conversation_id: Uuid,
before: Option<DateTime<Utc>>,
limit: i64,
) -> Result<(Vec<ConversationMessage>, bool), DatabaseError> {
self.inner
.list_conversation_messages_paginated(conversation_id, before, limit)
.await
}
pub async fn create_conversation_with_metadata(
&self,
channel: &str,
metadata: &serde_json::Value,
) -> Result<Uuid, DatabaseError> {
self.inner
.create_conversation_with_metadata(channel, &self.user_id, metadata)
.await
}
pub async fn update_conversation_metadata_field(
&self,
id: Uuid,
key: &str,
value: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.inner
.update_conversation_metadata_field(id, key, value)
.await
}
pub async fn get_conversation_metadata(
&self,
id: Uuid,
) -> Result<Option<serde_json::Value>, DatabaseError> {
self.inner.get_conversation_metadata(id).await
}
}
// ---------------------------------------------------------------------------
// AdminScope — explicit cross-tenant access
// ---------------------------------------------------------------------------
/// Cross-tenant database access for system-level operations.
///
/// **Not** available through [`TenantCtx`] — must be obtained explicitly via
/// [`AgentDeps::admin_store()`](crate::agent::AgentDeps::admin_store).
///
/// Used by: heartbeat enumeration, routine engine scheduling, self-repair,
/// scheduler job persistence, worker status updates.
#[derive(Clone)]
pub struct AdminScope {
inner: Arc<dyn Database>,
}
impl AdminScope {
pub fn new(db: Arc<dyn Database>) -> Self {
Self { inner: db }
}
/// Access the raw Database trait object.
///
/// Prefer using the typed methods on AdminScope instead. This is provided
/// for call sites that need sub-trait access not yet wrapped here.
pub fn db(&self) -> &Arc<dyn Database> {
&self.inner
}
// === Routine engine ===
pub async fn list_all_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
self.inner.list_all_routines().await
}
pub async fn list_event_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
self.inner.list_event_routines().await
}
pub async fn list_due_cron_routines(&self) -> Result<Vec<Routine>, DatabaseError> {
self.inner.list_due_cron_routines().await
}
pub async fn list_dispatched_routine_runs(&self) -> Result<Vec<RoutineRun>, DatabaseError> {
self.inner.list_dispatched_routine_runs().await
}
pub async fn count_running_routine_runs_batch(
&self,
routine_ids: &[Uuid],
) -> Result<HashMap<Uuid, i64>, DatabaseError> {
self.inner
.count_running_routine_runs_batch(routine_ids)
.await
}
pub async fn batch_get_last_run_status(
&self,
routine_ids: &[Uuid],
) -> Result<HashMap<Uuid, RunStatus>, DatabaseError> {
self.inner.batch_get_last_run_status(routine_ids).await
}
pub async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError> {
self.inner.count_running_routine_runs(routine_id).await
}
pub async fn update_routine_runtime(
&self,
id: Uuid,
last_run_at: DateTime<Utc>,
next_fire_at: Option<DateTime<Utc>>,
run_count: u64,
consecutive_failures: u32,
state: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.inner
.update_routine_runtime(
id,
last_run_at,
next_fire_at,
run_count,
consecutive_failures,
state,
)
.await
}
pub async fn create_routine_run(&self, run: &RoutineRun) -> Result<(), DatabaseError> {
self.inner.create_routine_run(run).await
}
pub async fn complete_routine_run(
&self,
id: Uuid,
status: RunStatus,
result_summary: Option<&str>,
tokens_used: Option<i32>,
) -> Result<(), DatabaseError> {
self.inner
.complete_routine_run(id, status, result_summary, tokens_used)
.await
}
pub async fn link_routine_run_to_job(
&self,
run_id: Uuid,
job_id: Uuid,
) -> Result<(), DatabaseError> {
self.inner.link_routine_run_to_job(run_id, job_id).await
}
pub async fn get_routine(&self, id: Uuid) -> Result<Option<Routine>, DatabaseError> {
self.inner.get_routine(id).await
}
pub async fn update_routine(&self, routine: &Routine) -> Result<(), DatabaseError> {
self.inner.update_routine(routine).await
}
// === Self-repair ===
pub async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError> {
self.inner.get_stuck_jobs().await
}
pub async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError> {
self.inner.get_broken_tools(threshold).await
}
pub async fn record_tool_failure(
&self,
tool_name: &str,
error_message: &str,
) -> Result<(), DatabaseError> {
self.inner
.record_tool_failure(tool_name, error_message)
.await
}
pub async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> {
self.inner.mark_tool_repaired(tool_name).await
}
pub async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> {
self.inner.increment_repair_attempts(tool_name).await
}
// === Sandbox housekeeping ===
pub async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError> {
self.inner.cleanup_stale_sandbox_jobs().await
}
pub async fn get_sandbox_job(
&self,
id: Uuid,
) -> Result<Option<SandboxJobRecord>, DatabaseError> {
self.inner.get_sandbox_job(id).await
}
pub async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
self.inner.save_sandbox_job(job).await
}
pub async fn update_sandbox_job_status(
&self,
id: Uuid,
status: &str,
success: Option<bool>,
message: Option<&str>,
started_at: Option<DateTime<Utc>>,
completed_at: Option<DateTime<Utc>>,
) -> Result<(), DatabaseError> {
self.inner
.update_sandbox_job_status(id, status, success, message, started_at, completed_at)
.await
}
pub async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> {
self.inner.update_sandbox_job_mode(id, mode).await
}
pub async fn get_sandbox_job_mode(&self, id: Uuid) -> Result<Option<String>, DatabaseError> {
self.inner.get_sandbox_job_mode(id).await
}
pub async fn save_job_event(
&self,
job_id: Uuid,
event_type: &str,
data: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.inner.save_job_event(job_id, event_type, data).await
}
pub async fn list_job_events(
&self,
job_id: Uuid,
limit: Option<i64>,
) -> Result<Vec<crate::history::JobEventRecord>, DatabaseError> {
self.inner.list_job_events(job_id, limit).await
}
// === Job persistence (scheduler, worker) ===
pub async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
self.inner.get_job(id).await
}
pub async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> {
self.inner.save_job(ctx).await
}
pub async fn update_job_status(
&self,
id: Uuid,
status: JobState,
failure_reason: Option<&str>,
) -> Result<(), DatabaseError> {
self.inner
.update_job_status(id, status, failure_reason)
.await
}
pub async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> {
self.inner.mark_job_stuck(id).await
}
pub async fn list_agent_jobs(&self) -> Result<Vec<AgentJobRecord>, DatabaseError> {
self.inner.list_agent_jobs().await
}
pub async fn get_agent_job_failure_reason(
&self,
id: Uuid,
) -> Result<Option<String>, DatabaseError> {
self.inner.get_agent_job_failure_reason(id).await
}
// === LLM call recording ===
pub async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError> {
self.inner.record_llm_call(record).await
}
pub async fn save_action(
&self,
job_id: Uuid,
action: &ActionRecord,
) -> Result<(), DatabaseError> {
self.inner.save_action(job_id, action).await
}
pub async fn get_job_actions(&self, job_id: Uuid) -> Result<Vec<ActionRecord>, DatabaseError> {
self.inner.get_job_actions(job_id).await
}
// === Estimation ===
pub async fn save_estimation_snapshot(
&self,
job_id: Uuid,
category: &str,
tool_names: &[String],
estimated_cost: Decimal,
estimated_time_secs: i32,
estimated_value: Decimal,
) -> Result<Uuid, DatabaseError> {
self.inner
.save_estimation_snapshot(
job_id,
category,
tool_names,
estimated_cost,
estimated_time_secs,
estimated_value,
)
.await
}
pub async fn update_estimation_actuals(
&self,
id: Uuid,
actual_cost: Decimal,
actual_time_secs: i32,
actual_value: Option<Decimal>,
) -> Result<(), DatabaseError> {
self.inner
.update_estimation_actuals(id, actual_cost, actual_time_secs, actual_value)
.await
}
// === Conversations (admin context) ===
pub async fn add_conversation_message(
&self,
conversation_id: Uuid,
role: &str,
content: &str,
) -> Result<Uuid, DatabaseError> {
self.inner
.add_conversation_message(conversation_id, role, content)
.await
}
pub async fn get_or_create_routine_conversation(
&self,
routine_id: Uuid,
routine_name: &str,
user_id: &str,
) -> Result<Uuid, DatabaseError> {
self.inner
.get_or_create_routine_conversation(routine_id, routine_name, user_id)
.await
}
pub async fn get_or_create_heartbeat_conversation(
&self,
user_id: &str,
) -> Result<Uuid, DatabaseError> {
self.inner
.get_or_create_heartbeat_conversation(user_id)
.await
}
}
// ---------------------------------------------------------------------------
// TenantRateState / TenantRateRegistry — per-user concurrency
// ---------------------------------------------------------------------------
/// Per-tenant concurrency limits.
pub struct TenantRateState {
/// Limits concurrent LLM calls for this user.
pub llm_semaphore: Arc<Semaphore>,
/// Limits concurrent jobs for this user.
pub job_semaphore: Arc<Semaphore>,
}
impl TenantRateState {
pub fn new(max_llm_concurrent: usize, max_job_concurrent: usize) -> Self {
Self {
llm_semaphore: Arc::new(Semaphore::new(max_llm_concurrent)),
job_semaphore: Arc::new(Semaphore::new(max_job_concurrent)),
}
}
}
/// Registry that lazily creates per-tenant rate state.
///
/// Uses `tokio::sync::RwLock<HashMap>` (consistent with the rest of the
/// codebase — no DashMap dependency).
pub struct TenantRateRegistry {
state: tokio::sync::RwLock<HashMap<String, Arc<TenantRateState>>>,
max_llm_concurrent: usize,
max_job_concurrent: usize,
}
impl TenantRateRegistry {
pub fn new(max_llm_concurrent: usize, max_job_concurrent: usize) -> Self {
Self {
state: tokio::sync::RwLock::new(HashMap::new()),
max_llm_concurrent,
max_job_concurrent,
}
}
/// Get or lazily create rate state for a user.
pub async fn get_or_create(&self, user_id: &str) -> Arc<TenantRateState> {
// Fast path: read lock
{
let map = self.state.read().await;
if let Some(s) = map.get(user_id) {
return Arc::clone(s);
}
}
// Slow path: write lock with double-check
let mut map = self.state.write().await;
if let Some(s) = map.get(user_id) {
return Arc::clone(s);
}
let s = Arc::new(TenantRateState::new(
self.max_llm_concurrent,
self.max_job_concurrent,
));
map.insert(user_id.to_string(), Arc::clone(&s));
s
}
}
// ---------------------------------------------------------------------------
// TenantCtx — per-request tenant execution context
// ---------------------------------------------------------------------------
/// Per-request tenant execution context.
///
/// Bundles a [`TenantScope`] (scoped DB access), workspace, cost guard,
/// and per-tenant rate limiting. Constructed once per request via
/// [`AgentDeps::tenant_ctx()`](crate::agent::AgentDeps::tenant_ctx).
///
/// `Clone + Send + Sync` — safe to store on `ChatDelegate` without lifetime issues.
#[derive(Clone)]
pub struct TenantCtx {
user_id: String,
store: Option<TenantScope>,
workspace: Option<Arc<Workspace>>,
cost_guard: Arc<CostGuard>,
rate: Arc<TenantRateState>,
}
impl TenantCtx {
pub fn new(
user_id: impl Into<String>,
store: Option<TenantScope>,
workspace: Option<Arc<Workspace>>,
cost_guard: Arc<CostGuard>,
rate: Arc<TenantRateState>,
) -> Self {
Self {
user_id: user_id.into(),
store,
workspace,
cost_guard,
rate,
}
}
pub fn user_id(&self) -> &str {
&self.user_id
}
pub fn store(&self) -> Option<&TenantScope> {
self.store.as_ref()
}
pub fn workspace(&self) -> Option<&Arc<Workspace>> {
self.workspace.as_ref()
}
pub fn cost_guard(&self) -> &CostGuard {
&self.cost_guard
}
/// Check cost limits for this tenant (global + per-user).
pub async fn check_cost_allowed(&self) -> Result<(), CostLimitExceeded> {
self.cost_guard.check_allowed_for_user(&self.user_id).await
}
/// Record an LLM call for this tenant.
#[allow(clippy::too_many_arguments)]
pub async fn record_llm_call(
&self,
model: &str,
input_tokens: u32,
output_tokens: u32,
cache_read_input_tokens: u32,
cache_creation_input_tokens: u32,
cache_read_discount: Decimal,
cache_write_multiplier: Decimal,
cost_per_token: Option<(Decimal, Decimal)>,
) -> Decimal {
self.cost_guard
.record_llm_call_for_user(
&self.user_id,
model,
input_tokens,
output_tokens,
cache_read_input_tokens,
cache_creation_input_tokens,
cache_read_discount,
cache_write_multiplier,
cost_per_token,
)
.await
}
/// Acquire an LLM concurrency permit for this tenant.
pub async fn acquire_llm_permit(&self) -> Result<SemaphorePermit<'_>, crate::error::Error> {
self.rate.llm_semaphore.acquire().await.map_err(|_| {
crate::error::Error::Config(crate::error::ConfigError::InvalidValue {
key: "llm_semaphore".to_string(),
message: "semaphore closed".to_string(),
})
})
}
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_rate_registry_returns_same_state_for_same_user() {
let registry = TenantRateRegistry::new(4, 3);
let a1 = registry.get_or_create("alice").await;
let a2 = registry.get_or_create("alice").await;
assert!(Arc::ptr_eq(&a1, &a2));
}
#[tokio::test]
async fn test_rate_registry_different_users_get_different_state() {
let registry = TenantRateRegistry::new(4, 3);
let alice = registry.get_or_create("alice").await;
let bob = registry.get_or_create("bob").await;
assert!(!Arc::ptr_eq(&alice, &bob));
}
}
+2
View File
@@ -532,6 +532,7 @@ impl TestHarnessBuilder {
let cost_guard = Arc::new(CostGuard::new(CostGuardConfig {
max_cost_per_day_cents: None,
max_actions_per_hour: None,
max_cost_per_user_per_day_cents: None,
}));
let channel = if self.stub_channel {
@@ -564,6 +565,7 @@ impl TestHarnessBuilder {
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
builder: None,
llm_backend: "nearai".to_string(),
tenant_rates: std::sync::Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
};
TestHarness {
+15 -1
View File
@@ -915,7 +915,7 @@ fn parse_routine_create_request(
fn build_routine_trigger(trigger: &NormalizedTriggerRequest) -> Trigger {
match trigger {
NormalizedTriggerRequest::Cron { schedule, timezone } => Trigger::Cron {
schedule: schedule.clone(),
schedule: normalize_cron_expression(schedule),
timezone: timezone.clone(),
},
NormalizedTriggerRequest::Manual => Trigger::Manual,
@@ -1836,6 +1836,20 @@ mod tests {
assert_eq!(parsed.cooldown_secs, 30);
}
#[test]
fn build_routine_trigger_normalizes_cron_schedule() {
let trigger = build_routine_trigger(&NormalizedTriggerRequest::Cron {
schedule: "0 0 9 * * MON-FRI".to_string(),
timezone: Some("UTC".to_string()),
});
assert!(matches!(
trigger,
Trigger::Cron { schedule, timezone }
if schedule == "0 0 9 * * MON-FRI *" && timezone.as_deref() == Some("UTC")
));
}
#[test]
fn parses_grouped_message_event_with_tools() {
let params = serde_json::json!({
+72 -4
View File
@@ -18,8 +18,8 @@ use crate::agent::agentic_loop::{
};
use crate::agent::scheduler::WorkerMessage;
use crate::agent::task::TaskOutput;
use crate::channels::web::types::ToolDecisionDto;
use crate::context::{ContextManager, JobState};
use crate::db::Database;
use crate::error::Error;
use crate::hooks::HookRegistry;
use crate::llm::{
@@ -27,6 +27,7 @@ use crate::llm::{
ToolSelection,
};
use crate::safety::SafetyLayer;
use crate::tenant::AdminScope;
use crate::tools::execute::process_tool_result;
use crate::tools::rate_limiter::RateLimitResult;
use crate::tools::{
@@ -44,7 +45,7 @@ pub struct WorkerDeps {
pub llm: Arc<dyn LlmProvider>,
pub safety: Arc<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
pub store: Option<Arc<dyn Database>>,
pub store: Option<AdminScope>,
pub hooks: Arc<HookRegistry>,
pub timeout: Duration,
pub use_planning: bool,
@@ -93,7 +94,7 @@ impl Worker {
&self.deps.tools
}
fn store(&self) -> Option<&Arc<dyn Database>> {
fn store(&self) -> Option<&AdminScope> {
self.deps.store.as_ref()
}
@@ -200,6 +201,19 @@ impl Worker {
.map(|s| s.to_string()),
fallback_deliverable: data.get("fallback_deliverable").cloned(),
}),
"reasoning" => {
let narrative = data
.get("narrative")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let decisions = ToolDecisionDto::from_json_array(&data["decisions"]);
Some(AppEvent::JobReasoning {
job_id: job_id_str,
narrative,
decisions,
})
}
_ => None,
};
if let Some(event) = event {
@@ -897,6 +911,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
id: selection.tool_call_id.clone(),
name: selection.tool_name.clone(),
arguments: selection.parameters.clone(),
reasoning: if action.reasoning.is_empty() {
None
} else {
Some(action.reasoning.clone())
},
}],
));
@@ -1139,6 +1158,7 @@ impl<'a> JobDelegate<'a> {
Ok(crate::llm::RespondOutput {
result: RespondResult::Text(String::new()),
usage: crate::llm::TokenUsage::default(),
finish_reason: crate::llm::FinishReason::Stop,
})
}
}
@@ -1264,6 +1284,7 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
content: reasoning_text,
},
usage: crate::llm::TokenUsage::default(),
finish_reason: crate::llm::FinishReason::ToolUse,
});
}
Ok(_) => {} // empty selections, fall through
@@ -1357,6 +1378,48 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
);
}
// Emit reasoning event if any tool calls carry reasoning.
// Sanitize narrative and per-tool rationale through SafetyLayer
// (parity with ChatDelegate in dispatcher.rs).
let sanitized_narrative = content
.as_deref()
.filter(|c| !c.trim().is_empty())
.map(|c| {
self.worker
.deps
.safety
.sanitize_tool_output("job_narrative", c)
.content
})
.filter(|c| !c.trim().is_empty())
.unwrap_or_default();
let decisions: Vec<serde_json::Value> = tool_calls
.iter()
.filter_map(|tc| {
tc.reasoning.as_ref().map(|r| {
let sanitized = self
.worker
.deps
.safety
.sanitize_tool_output("tool_rationale", r)
.content;
serde_json::json!({
"tool_name": tc.name,
"rationale": sanitized,
})
})
})
.collect();
if !decisions.is_empty() {
self.worker.log_event(
"reasoning",
serde_json::json!({
"narrative": sanitized_narrative,
"decisions": decisions,
}),
);
}
// Add assistant message with tool_calls (OpenAI protocol)
reason_ctx
.messages
@@ -1371,7 +1434,7 @@ impl<'a> LoopDelegate for JobDelegate<'a> {
.map(|tc| ToolSelection {
tool_name: tc.name.clone(),
parameters: tc.arguments.clone(),
reasoning: String::new(),
reasoning: tc.reasoning.clone().unwrap_or_default(),
alternatives: vec![],
tool_call_id: tc.id.clone(),
})
@@ -1424,6 +1487,11 @@ fn selections_to_tool_calls(selections: &[ToolSelection]) -> Vec<ToolCall> {
id: s.tool_call_id.clone(),
name: s.tool_name.clone(),
arguments: s.parameters.clone(),
reasoning: if s.reasoning.is_empty() {
None
} else {
Some(s.reasoning.clone())
},
})
.collect()
}
+55
View File
@@ -149,6 +149,7 @@ fn reject_if_injected(path: &str, content: &str) -> Result<(), WorkspaceError> {
///
/// Allows Workspace to work with either a PostgreSQL `Repository` (the original
/// path) or any `Database` trait implementation (e.g. libSQL backend).
#[derive(Clone)]
enum WorkspaceStorage {
/// PostgreSQL-backed repository (uses connection pool directly).
#[cfg(feature = "postgres")]
@@ -576,6 +577,60 @@ impl Workspace {
self
}
/// Clone the workspace configuration for a different primary user scope.
///
/// This preserves search config, embeddings, shared read scopes, memory
/// layers, and privacy classifier while switching the primary read/write
/// scope to `user_id`.
pub fn scoped_to_user(&self, user_id: impl Into<String>) -> Self {
let user_id = user_id.into();
let mut memory_layers = self.memory_layers.clone();
for layer in &mut memory_layers {
if layer.sensitivity == crate::workspace::layer::LayerSensitivity::Private
&& layer.scope == self.user_id
{
layer.scope = user_id.clone();
}
}
let mut read_user_ids = vec![user_id.clone()];
for scope in &self.read_user_ids {
if scope != &self.user_id && !read_user_ids.contains(scope) {
read_user_ids.push(scope.clone());
}
}
for scope in crate::workspace::layer::MemoryLayer::read_scopes(&memory_layers) {
if !read_user_ids.contains(&scope) {
read_user_ids.push(scope);
}
}
let preserve_flags = user_id == self.user_id;
Self {
user_id,
read_user_ids,
agent_id: self.agent_id,
storage: self.storage.clone(),
embeddings: self.embeddings.clone(),
bootstrap_pending: std::sync::atomic::AtomicBool::new(if preserve_flags {
self.bootstrap_pending
.load(std::sync::atomic::Ordering::Acquire)
} else {
false
}),
bootstrap_completed: std::sync::atomic::AtomicBool::new(if preserve_flags {
self.bootstrap_completed
.load(std::sync::atomic::Ordering::Acquire)
} else {
false
}),
search_defaults: self.search_defaults.clone(),
memory_layers,
privacy_classifier: self.privacy_classifier.clone(),
}
}
/// Get the user ID (primary scope for writes).
pub fn user_id(&self) -> &str {
&self.user_id
+1
View File
@@ -15,6 +15,7 @@ use crate::workspace::document::{MemoryChunk, MemoryDocument, WorkspaceEntry};
use crate::workspace::search::{RankedResult, SearchConfig, SearchResult, fuse_results};
/// Database repository for workspace operations.
#[derive(Clone)]
pub struct Repository {
pool: Pool,
}
@@ -253,6 +253,6 @@ async def test_telegram_hot_activation_transitions_installed_to_active(page):
assert await card.locator(SEL["ext_pairing_label"]).count() == 0
assert captured_setup_payloads == [
{"secrets": {"telegram_bot_token": "123456789:ABCdefGhI"}},
{"secrets": {}},
{"secrets": {"telegram_bot_token": "123456789:ABCdefGhI"}, "fields": {}},
{"secrets": {}, "fields": {}},
]
+3 -2
View File
@@ -587,6 +587,7 @@ mod advanced {
async fn mcp_extension_lifecycle() {
use crate::support::mock_mcp_server::{MockToolResponse, start_mock_mcp_server};
use ironclaw::extensions::{AuthHint, ExtensionKind, ExtensionSource, RegistryEntry};
const TEST_USER_ID: &str = "test-user";
// 1. Start mock MCP server with pre-configured tool responses.
let mock_server = start_mock_mcp_server(vec![
@@ -654,14 +655,14 @@ mod advanced {
ext_mgr
.secrets()
.create(
"default",
TEST_USER_ID,
ironclaw::secrets::CreateSecretParams::new(secret_name, "mock-access-token")
.with_provider("mcp:mock-notion".to_string()),
)
.await
.expect("failed to inject test token");
let activate_result = ext_mgr.activate("mock-notion", "default").await;
let activate_result = ext_mgr.activate("mock-notion", TEST_USER_ID).await;
assert!(
activate_result.is_ok(),
"activation failed: {:?}",
+1 -1
View File
@@ -439,7 +439,7 @@ mod tests {
match &routine.trigger {
Trigger::Cron { schedule, timezone } => {
assert_eq!(schedule, "0 0 9 * * MON-FRI");
assert_eq!(schedule, "0 0 9 * * MON-FRI *");
assert_eq!(timezone.as_deref(), Some("UTC"));
}
other => panic!("expected cron trigger, got {other:?}"),
+10 -10
View File
@@ -337,14 +337,14 @@ mod tests {
SchedulerDeps {
tools: registry.clone(),
extension_manager: extension_manager.clone(),
store: Some(db.clone()),
store: Some(ironclaw::tenant::AdminScope::new(db.clone())),
hooks: Arc::new(HookRegistry::new()),
},
));
Arc::new(RoutineEngine::new(
RoutineConfig::default(),
db,
ironclaw::tenant::AdminScope::new(db),
llm,
ws,
notify_tx,
@@ -448,7 +448,7 @@ mod tests {
let engine = Arc::new(RoutineEngine::new(
RoutineConfig::default(),
db.clone(),
ironclaw::tenant::AdminScope::new(db.clone()),
llm,
ws,
notify_tx,
@@ -527,7 +527,7 @@ mod tests {
let engine = Arc::new(RoutineEngine::new(
RoutineConfig::default(),
db.clone(),
ironclaw::tenant::AdminScope::new(db.clone()),
llm,
ws,
notify_tx,
@@ -614,7 +614,7 @@ mod tests {
let engine = Arc::new(RoutineEngine::new(
RoutineConfig::default(),
db.clone(),
ironclaw::tenant::AdminScope::new(db.clone()),
llm,
ws,
notify_tx,
@@ -723,7 +723,7 @@ mod tests {
let engine = Arc::new(RoutineEngine::new(
RoutineConfig::default(),
db.clone(),
ironclaw::tenant::AdminScope::new(db.clone()),
llm,
ws,
notify_tx,
@@ -866,7 +866,7 @@ mod tests {
let engine = Arc::new(RoutineEngine::new(
RoutineConfig::default(),
db.clone(),
ironclaw::tenant::AdminScope::new(db.clone()),
llm,
ws,
notify_tx,
@@ -1049,7 +1049,7 @@ mod tests {
let engine = Arc::new(RoutineEngine::new(
RoutineConfig::default(),
Arc::clone(&db),
ironclaw::tenant::AdminScope::new(Arc::clone(&db)),
llm,
ws,
notify_tx,
@@ -1171,7 +1171,7 @@ mod tests {
let engine = Arc::new(RoutineEngine::new(
RoutineConfig::default(),
db.clone(),
ironclaw::tenant::AdminScope::new(db.clone()),
llm,
ws,
notify_tx,
@@ -1279,7 +1279,7 @@ mod tests {
let engine = Arc::new(RoutineEngine::new(
config,
db.clone(),
ironclaw::tenant::AdminScope::new(db.clone()),
llm,
ws,
notify_tx,
+1
View File
@@ -201,6 +201,7 @@ mod tests {
sandbox_readiness: ironclaw::agent::SandboxReadiness::DisabledByConfig,
builder: None,
llm_backend: "nearai".to_string(),
tenant_rates: std::sync::Arc::new(ironclaw::tenant::TenantRateRegistry::new(4, 3)),
};
let gateway = Arc::new(TestChannel::new());
+3 -1
View File
@@ -12,6 +12,7 @@ mod tests {
use crate::support::test_rig::TestRigBuilder;
use crate::support::trace_llm::LlmTrace;
use ironclaw::workspace::Workspace;
// -----------------------------------------------------------------------
// Test 1: write_chunk_search
@@ -268,6 +269,7 @@ mod tests {
#[tokio::test]
async fn identity_in_system_prompt() {
const TEST_USER_ID: &str = "test-user";
let trace = LlmTrace::from_file(concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/llm_traces/workspace/identity_prompt.json"
@@ -280,7 +282,7 @@ mod tests {
.await;
// Seed an IDENTITY.md so the system prompt has real content to inject.
let ws = rig.workspace().expect("workspace must be available");
let ws = Workspace::new_with_db(TEST_USER_ID, rig.database().clone());
ws.write(
"IDENTITY.md",
"I am TestBot, a helpful testing assistant created for E2E verification.",
+7 -7
View File
@@ -1,10 +1,10 @@
//! Tests proving that multi-tenant system prompts are broken.
//! Regression tests for multi-tenant system prompts.
//!
//! Bug: In multi-tenant mode, the agent loop uses `self.workspace()` which
//! returns a single shared workspace (user_id="default"). Identity files
//! (IDENTITY.md, SOUL.md, USER.md) seeded under per-user IDs ("alice",
//! "bob") are invisible to this workspace, so the system prompt is
//! empty/wrong.
//! The agent must build the conversational system prompt from a workspace
//! scoped to the incoming message's user, not from the shared owner-scope
//! workspace created at startup. Otherwise per-user identity files
//! (IDENTITY.md, SOUL.md, USER.md) become invisible and different users can
//! see the same owner-scoped prompt.
//!
//! These tests:
//! 1. Seed identity files for two users (alice, bob) in the database
@@ -13,7 +13,7 @@
//! correct user's identity
//! 4. Verify user A's identity doesn't leak into user B's prompt
//!
//! All tests are expected to FAIL until the bug is fixed.
//! These tests ensure each user's identity is isolated correctly.
#[cfg(feature = "libsql")]
mod support;
+1
View File
@@ -94,6 +94,7 @@ impl LlmProvider for MockLlmProvider {
id: "call_mock_001".to_string(),
name: tool.name.clone(),
arguments: serde_json::json!({"test": true}),
reasoning: None,
}],
input_tokens: 15,
output_tokens: 8,
@@ -267,6 +267,7 @@ impl GatewayWorkflowHarness {
sandbox_readiness: ironclaw::agent::SandboxReadiness::DisabledByConfig,
builder: None,
llm_backend: "nearai".to_string(),
tenant_rates: std::sync::Arc::new(ironclaw::tenant::TenantRateRegistry::new(4, 3)),
},
channels,
None,
+2 -1
View File
@@ -642,7 +642,7 @@ impl TestRigBuilder {
let (notify_tx, _notify_rx) = tokio::sync::mpsc::channel(16);
let engine = Arc::new(RoutineEngine::new(
routine_config,
Arc::clone(db_arc),
ironclaw::tenant::AdminScope::new(Arc::clone(db_arc)),
components.llm.clone(),
Arc::clone(ws),
notify_tx,
@@ -762,6 +762,7 @@ impl TestRigBuilder {
sandbox_readiness: ironclaw::agent::SandboxReadiness::Available, // tests don't use real Docker
builder: None,
llm_backend: "nearai".to_string(),
tenant_rates: std::sync::Arc::new(ironclaw::tenant::TenantRateRegistry::new(4, 3)),
};
// 7. Create TestChannel and ChannelManager.
+1
View File
@@ -566,6 +566,7 @@ impl LlmProvider for TraceLlm {
id: tc.id,
name: tc.name,
arguments: tc.arguments,
reasoning: None,
})
.collect();
Ok(ToolCompletionResponse {