feat(users): UserManager with per-user SQLCipher, and extract skald-core crate

Two changes developed together in one session; they share the same module
structure (db/mod.rs, the core lib root) and only compile together, so they
land as one commit.

## UserManager + per-user encryption (§9/§11)

New `users::UserManager`: owns the system.db pool plus a map
`userid -> SqlitePool` of unlocked databases. The pool *is* the unlock token —
its connect options carry the DEK as SQLCipher's raw key, so an open pool means
the key is in RAM until restart and dropping it re-locks (§9). Knows nothing
about cookies.

New `crypto` module: envelope encryption. A random 256-bit DEK encrypts
`{userid}.db`; `users.database_password` holds it sealed with AES-256-GCM under
`Argon2id(password, salt)`. The AEAD tag is the password verifier — one
derivation both authenticates and yields the key, so encrypted users store no
second hash. Cleartext users store the Argon2id output directly, compared in
constant time. Argon2 runs in spawn_blocking behind a 2-permit semaphore
(256 MiB per derivation).

- SQLCipher via `libsqlite3-sys` `bundled-sqlcipher-vendored-openssl`, pinned
  <0.38 so it unifies with the one sqlx-sqlite links (a newer copy would apply
  the feature to a SQLite sqlx never uses). OpenSSL is vendored and static, so
  the binary stays self-contained.
- Schema split into `create_registry_tables` (instance-wide, no user key) and
  `create_owner_tables` (one owner's content, identical in every file). No FK in
  the owner bucket may reach the registry — enforced by a standalone test.
  Dropped `chat_history.model_db_id` (write-only, and the only registry-crossing
  key); moved `projects`/`project_tickets` into the owner bucket.
- Provisioning invariant: the file is written before the row, deleted after it,
  so a crash leaves an orphan file, never a user without a database. `open_db`
  never creates: a missing file is an error, not a silent empty database.

Not consumed yet: no login, call sites still use the shared system.db pool.

## Extract crates/skald-core

The headless core moves out of `src/` into its own crate; `skald` (server) and
the coming `skald-setup` are shells around it. Two dependencies on the shell
were inverted rather than dragged along, so the core names neither Tauri nor any
concrete plugin:

- `Plugin::tools(self: Arc<Self>)` — plugins contribute tools through this hook
  (sibling of `http_router`), so the core no longer downcasts to
  `MobileConnectorPlugin`.
- `tools::restart::set_restart_handler` — the desktop shell installs its
  teardown-and-respawn; the core defaults to the supervisor exit code. The core
  loses its `desktop` feature.
- `boot`'s stdout formatter moves to the binary (`src/boot_format.rs`); the core
  only emits tracing events.

All 79 core tests pass; the binary boots and serves in a clean directory, and
the mobile-connector tools still register through the new hook.
This commit is contained in:
2026-07-10 16:48:51 +01:00
parent 38494a85a9
commit 178a38357e
173 changed files with 2650 additions and 1106 deletions
+82
View File
@@ -0,0 +1,82 @@
[package]
name = "skald-core"
version = "0.1.0"
edition = "2024"
# The headless application core: database + crypto + identity, the LLM stack,
# tools, MCP, plugins-as-a-registry, sessions.
#
# Deliberately knows nothing about the process shell around it. No Tauri, no
# concrete plugin crates (it only ever sees `Arc<dyn Plugin>` from `core-api`),
# no `axum` server — `skald` and `skald-setup` are both consumers.
[dependencies]
axum = { version = "0.8", features = ["ws", "multipart"] }
tokio = { version = "1.52.3", features = ["full"] }
tokio-util = { version = "0.7", features = ["rt"] }
futures = "0.3"
serde = { version = "1", features = ["derive"] }
anyhow = "1"
sqlx = { version = "0.9.0", features = ["runtime-tokio", "sqlite"] }
# Per-user database encryption. `sqlx-sqlite` already builds `libsqlite3-sys`;
# naming it here only adds a feature, and Cargo's feature unification turns the
# one SQLite the process links into SQLCipher. A database opened without
# `PRAGMA key` still behaves as plain SQLite, so `system.db` is unaffected.
# `-vendored-openssl` compiles libcrypto from source instead of linking the
# system one, so the binary stays self-contained (see b08a13c). It costs a C
# build (perl + make) — the trade the blueprint §4 accepts explicitly.
#
# Pinned below 0.38 on purpose: `sqlx-sqlite` 0.9 requires ">=0.30.1, <0.38.0".
# A newer release would resolve to a *second* copy of the crate, and the
# SQLCipher feature would then apply to a SQLite that sqlx never links. 0.37.0
# is the newest version that still unifies.
libsqlite3-sys = { version = "0.37", features = ["bundled-sqlcipher-vendored-openssl"] }
# 0.6 is still a release candidate; stay on the stable line.
argon2 = "0.5"
# Held at 0.10 to match `skald-relay-common`, so the binary links one AES-GCM
# rather than two. Bumping to 0.11 means bumping the mobile E2E code with it.
aes-gcm = "0.10"
zeroize = { version = "1", features = ["derive"] }
subtle = "2"
uuid = { version = "1", features = ["v4"] }
reqwest = { version = "0.13.4", default-features = false, features = ["rustls-no-provider", "charset", "http2", "system-proxy", "json", "multipart"] }
async-trait = "0.1"
serde_json = "1"
indexmap = { version = "2", features = ["serde"] }
tracing = "0.1"
chrono = { version = "0.4", default-features = false, features = ["clock", "std"] }
chrono-tz = "0.10"
iana-time-zone = "0.1"
os_info = "3"
cron = "0.16"
rand = "0.10.1"
syn = { version = "2", features = ["full"] }
quote = "1"
proc-macro2 = { version = "1", features = ["span-locations"] }
tree-sitter = "0.26"
tree-sitter-python = "0.25"
tree-sitter-javascript = "0.25"
tree-sitter-typescript = "0.23"
tree-sitter-go = "0.25"
tree-sitter-java = "0.23"
tree-sitter-c = "0.24"
tree-sitter-cpp = "0.23"
tree-sitter-swift = "0.7"
tree-sitter-lua = "0.5"
tree-sitter-ruby = "0.23"
tree-sitter-bash = "0.25"
tree-sitter-elixir = "0.3"
tree-sitter-json = "0.24"
tree-sitter-yaml = "0.7"
tree-sitter-html = "0.23"
tree-sitter-css = "0.25"
regex = "1"
glob = "0.3"
libc = "0.2"
base64 = "0.22"
sha2 = "0.10"
notify = "8"
honcho-client = { path = "../honcho-client" }
llm-client = { path = "../llm-client" }
core-api = { path = "../core-api" }
mcp-client = { path = "../mcp-client" }
+259
View File
@@ -0,0 +1,259 @@
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use tracing::{debug, trace, warn};
use core_api::provider::LlmStrength;
const AGENTS_DIR: &str = "agents";
/// The role an agent plays, declared by the required `type` field in `meta.json`.
///
/// - `Chat`: a conversational entry-point the user talks to directly (e.g. `main`,
/// `project-coordinator`). Not dispatchable as a sub-agent, not a valid task root.
/// - `Task`: a task executor. Dispatchable by a parent agent **and** a valid root of a
/// scheduled/async task (e.g. `software-engineer`, `researcher`, `generalist`).
/// - `System`: a hidden background agent wired into the runtime by id (e.g. `tic`).
/// Never listed, never user-chattable, never dispatchable from the tool surface.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AgentType {
Chat,
Task,
System,
}
#[derive(Deserialize)]
struct RawMeta {
name: String,
description: String,
#[serde(default)]
friendly_description: Option<String>,
#[serde(default)]
instructions: Option<String>,
#[serde(default)]
inject_memory: Vec<String>,
#[serde(default)]
client: Option<String>,
#[serde(default)]
scope: Option<String>,
#[serde(default)]
strength: Option<LlmStrength>,
/// Required: declares the agent's role. A `meta.json` without `type` fails to load.
#[serde(rename = "type")]
agent_type: AgentType,
#[serde(default = "default_true")]
inject_skills: bool,
#[serde(default)]
icon: Option<String>,
}
/// Serde default for boolean fields that should be `true` when the key is absent.
fn default_true() -> bool { true }
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentMeta {
pub id: String,
pub name: String,
/// Routing description for the **orchestrator LLM**: "when should I delegate to this
/// agent, and what does it return?". Injected into `<!-- AGENTS_LIST -->` and returned
/// by `list_items` (type=agents). Required.
pub description: String,
/// Human-facing blurb shown to the **user** on the frontend Agents page. When absent
/// the frontend falls back to `description`. Applies to every agent type.
#[serde(default)]
pub friendly_description: Option<String>,
/// Note for the **calling LLM** on *how* to invoke this agent for the best result
/// (expected inputs, format, gotchas). Kept short. Only meaningful for `task` agents —
/// it is surfaced solely via `list_items` (type=agents), which already lists task
/// agents only, so no extra gating is needed.
#[serde(default)]
pub instructions: Option<String>,
#[serde(default)]
pub inject_memory: Vec<String>,
/// Preferred LLM client name (must exist in the DB, configured via the web app).
/// If unset, the sub-agent inherits the caller's client.
#[serde(default)]
pub client: Option<String>,
/// Task domain this agent operates in (e.g. "coding", "reasoning").
/// Used by AUTO client selection to find a matching LLM.
#[serde(default)]
pub scope: Option<String>,
/// Minimum LLM capability required to run this agent reliably.
/// AUTO selection skips clients weaker than this threshold.
#[serde(default)]
pub strength: Option<LlmStrength>,
/// The agent's role (`chat` / `task` / `system`). Only `task` agents are listed in
/// `list_items` (type=agents) / the AGENTS_LIST injection and are dispatchable or
/// runnable as a task root; `chat` and `system` are excluded from those paths.
#[serde(rename = "type")]
pub agent_type: AgentType,
/// When true (the default, including when the key is absent), the skills index
/// (`skills/index.md`) is injected into this agent's system prompt so it can
/// discover and use installed skills. Set false for background agents that don't
/// need them (e.g. TIC) to save tokens.
#[serde(default = "default_true")]
pub inject_skills: bool,
/// Path to the agent's icon image file (relative to the agent's directory).
/// Defaults to None if no icon is configured.
#[serde(default)]
pub icon: Option<String>,
}
/// Scan `agents/` and return metadata for every agent that has both
/// `meta.json` and `AGENT.md`. Skips the `common/` directory.
pub fn discover() -> Result<Vec<AgentMeta>> {
let mut agents = Vec::new();
let dir = std::fs::read_dir(AGENTS_DIR)
.with_context(|| format!("Failed to read agents directory '{AGENTS_DIR}'"))?;
for entry in dir {
let entry = entry?;
let path = entry.path();
if !path.is_dir() { continue; }
let id = match path.file_name().and_then(|n| n.to_str()) {
Some(n) if !n.is_empty() && n != "common" => n.to_string(),
_ => continue,
};
let meta_path = path.join("meta.json");
let system_path = path.join("AGENT.md");
if !meta_path.exists() || !system_path.exists() {
warn!(agent_id = %id, "skipping agent: missing meta.json or AGENT.md");
continue;
}
let raw_str = match std::fs::read_to_string(&meta_path) {
Ok(s) => s,
Err(e) => {
warn!(agent_id = %id, error = %e, "skipping agent: cannot read meta.json");
continue;
}
};
// A single malformed meta.json (e.g. missing the required `type` field) must not
// blank the whole roster — warn and skip it, keep discovering the rest.
let raw: RawMeta = match serde_json::from_str(&raw_str) {
Ok(r) => r,
Err(e) => {
warn!(agent_id = %id, error = %e, "skipping agent: invalid meta.json");
continue;
}
};
let meta = AgentMeta {
id,
name: raw.name,
description: raw.description,
friendly_description: raw.friendly_description,
instructions: raw.instructions,
inject_memory: raw.inject_memory,
client: raw.client,
scope: raw.scope,
strength: raw.strength,
agent_type: raw.agent_type,
inject_skills: raw.inject_skills,
icon: raw.icon,
};
trace!(agent_id = %meta.id, client = ?meta.client, scope = ?meta.scope, strength = ?meta.strength, "agent meta loaded");
debug!(agent_id = %meta.id, name = %meta.name, "agent discovered");
agents.push(meta);
}
agents.sort_by(|a, b| a.id.cmp(&b.id));
Ok(agents)
}
/// Load metadata for a single agent (reads its `meta.json`).
pub fn load_meta(agent_id: &str) -> Result<AgentMeta> {
let path = format!("{AGENTS_DIR}/{agent_id}/meta.json");
let raw_str = std::fs::read_to_string(&path)
.with_context(|| format!("Agent '{agent_id}': meta.json not found at '{path}'"))?;
let raw: RawMeta = serde_json::from_str(&raw_str)
.with_context(|| format!("Agent '{agent_id}': failed to parse meta.json"))?;
Ok(AgentMeta {
id: agent_id.to_string(),
name: raw.name,
description: raw.description,
friendly_description: raw.friendly_description,
instructions: raw.instructions,
inject_memory: raw.inject_memory,
client: raw.client,
scope: raw.scope,
strength: raw.strength,
agent_type: raw.agent_type,
inject_skills: raw.inject_skills,
icon: raw.icon,
})
}
/// Load metadata for `agent_id` and assert it is a runnable **task** agent.
/// Errors if the agent does not exist or is a `chat` / `system` agent — i.e. the
/// single gate for "can this agent be dispatched or run as a task root?".
pub fn load_task_meta(agent_id: &str) -> Result<AgentMeta> {
let meta = load_meta(agent_id)?;
if meta.agent_type != AgentType::Task {
anyhow::bail!(
"agent `{agent_id}` is a {:?} agent and cannot be dispatched or run as a task — only `task` agents can",
meta.agent_type
);
}
Ok(meta)
}
/// Load and resolve the system prompt for `agent_id` from disk.
/// Called at request time so edits to `.md` files take effect without restart.
pub fn load_prompt(agent_id: &str) -> Result<String> {
let path = format!("{AGENTS_DIR}/{agent_id}/AGENT.md");
let content = std::fs::read_to_string(&path)
.with_context(|| format!("Agent '{agent_id}': AGENT.md not found at '{path}'"))?;
resolve_includes(&content)
}
fn resolve_includes(content: &str) -> Result<String> {
let mut out = String::with_capacity(content.len());
for line in content.lines() {
let trimmed = line.trim();
if let Some(path_raw) = trimmed
.strip_prefix("<!-- INCLUDE:")
.and_then(|s| s.strip_suffix("-->"))
{
let path = format!("{AGENTS_DIR}/{}", path_raw.trim());
let included = std::fs::read_to_string(&path)
.with_context(|| format!("INCLUDE: failed to read '{path}'"))?;
out.push_str(&format!("<included_file path=\"{path}\">\n"));
out.push_str(&resolve_includes(&included)?);
out.push_str("</included_file>\n");
} else if trimmed == "<!-- AGENTS_LIST -->" {
out.push_str(&render_agents_list()?);
} else if trimmed == "<!-- MCP_LIST -->" {
// Replaced at request time in build_openai_messages with dynamic
// active/hidden sections. Leave a sentinel so the injection point
// is preserved and positioned correctly in the prompt.
out.push_str("__MCP_LIST__\n");
} else if let Some(key) = trimmed
.strip_prefix("<!-- ")
.and_then(|s| s.strip_suffix(" -->"))
.filter(|k| k.chars().all(|c| c.is_ascii_uppercase() || c == '_'))
{
// Generic runtime substitution: <!-- KEY --> → __KEY__ sentinel.
// Replaced at request time via SendMessageOptions::system_substitutions.
out.push_str(&format!("__{key}__\n"));
} else {
out.push_str(line);
out.push('\n');
}
}
Ok(out)
}
fn render_agents_list() -> Result<String> {
let agents = discover()?;
let mut out = String::new();
for agent in agents.iter().filter(|a| a.agent_type == AgentType::Task) {
out.push_str(&format!("- **{}** — {}\n", agent.id, agent.description));
}
Ok(out)
}
File diff suppressed because it is too large Load Diff
+53
View File
@@ -0,0 +1,53 @@
//! Curated, human-readable bootstrap progress printed to **stdout**.
//!
//! At runtime stdout is silent — only the file log (`logs/skald.log`) records
//! events. During startup, though, it is useful to see at a glance how the app
//! is configured and how it is coming up. These helpers emit a small, ordered
//! set of lines on the dedicated `boot` tracing target.
//!
//! Rendering belongs to whoever owns the process: each shell installs a stdout
//! layer filtered on [`TARGET`] and formats the lines as it likes (`skald`'s
//! `BootFormat` strips timestamps and paints failures red). The core only says
//! *what* happened, never how it looks — which is also why nothing here depends
//! on `tracing-subscriber`.
//!
//! The same lines land in the log file (they pass the normal `EnvFilter`), so
//! they double as a high-level startup trace. Glyphs are baked into the message
//! on purpose, so the file keeps the same readable shape.
use std::fmt;
use tracing::{info, warn};
/// Tracing target for curated bootstrap lines shown on stdout.
pub const TARGET: &str = "boot";
/// Top-level title (no glyph), e.g. `skald v0.5 — starting`.
pub fn title(msg: impl fmt::Display) {
info!(target: TARGET, "{}", msg);
}
/// A phase header, e.g. ` Plugins — 6 active, 1 failed`.
pub fn section(msg: impl fmt::Display) {
info!(target: TARGET, " {}", msg);
}
/// A successful item under a phase.
pub fn ok(msg: impl fmt::Display) {
info!(target: TARGET, " ✓ {}", msg);
}
/// An item that exists but is inactive (e.g. a disabled plugin).
pub fn off(msg: impl fmt::Display) {
info!(target: TARGET, " ○ {}", msg);
}
/// A failed item (rendered in red on stdout; logged at WARN in the file).
pub fn fail(msg: impl fmt::Display) {
warn!(target: TARGET, " ✗ {}", msg);
}
/// The final "app is up" line, e.g. `✅ Ready — http://localhost:8080`.
pub fn ready(msg: impl fmt::Display) {
info!(target: TARGET, "✅ {}", msg);
}
+3
View File
@@ -0,0 +1,3 @@
pub use core_api::bus::{
BusEvent, ChatEvent, ChatEventBus, ChatEventRole, CompactionEvent, RecvError, ToolCallEvent,
};
+167
View File
@@ -0,0 +1,167 @@
//! Per-source input inbox for ChatHub.
//!
//! Each interactive source (telegram, web, mobile…) gets one `SourceInbox` and a
//! single consumer task (spawned lazily in `ChatHub`). A single consumer per
//! source makes delivery strictly FIFO, removing the ordering race of the old
//! detached-spawn dispatch.
//!
//! Messages are kept as **individual** units — they are not coalesced here. The
//! consumer pops one to seed a turn (`build_unit`); any further messages that
//! pile up while the turn runs are drained, one row each, at the turn's round
//! boundaries (`drain_leading_user`) and injected live into the running turn.
//! Coalescing for the LLM (merging consecutive user rows into one `role:user`)
//! happens later in the `MessageBuilder`, not here, so the DB keeps each message
//! distinct while the model still sees a single clean user turn.
//!
//! Serialization of the turns themselves still lives in
//! `ChatSessionHandler.processing`; this inbox sits in front of it, adding ordering.
use std::collections::VecDeque;
use std::sync::atomic::AtomicU64;
use tokio::sync::{Mutex, Notify};
use core_api::chat_hub::SendMessageOptions;
use core_api::message_meta::MessageMetadata;
/// One queued user message awaiting dispatch.
pub(super) struct QueuedMessage {
pub prompt: String,
pub opts: SendMessageOptions,
}
/// Pending queue + wake signal for a single source.
#[derive(Default)]
pub(super) struct SourceInbox {
pub pending: Mutex<VecDeque<QueuedMessage>>,
pub notify: Notify,
/// Bumped by `ChatHub::cancel` (after clearing `pending`) so the consumer can
/// drop a unit it drained microseconds before a `/stop`.
pub cancel_epoch: AtomicU64,
}
/// Pops the next dispatch unit from `pending` — a **single** message, used by the
/// consumer to seed a turn. No coalescing: any further queued messages are drained
/// into the running turn at its round boundaries (see `drain_leading_user`).
///
/// Empty queue → `None`. Synthetic messages (notification/TIC) and plain user
/// messages are treated identically here; only `drain_leading_user` distinguishes
/// them, leaving synthetic ones for the notification path.
pub(super) fn build_unit(
pending: &mut VecDeque<QueuedMessage>,
) -> Option<(String, SendMessageOptions)> {
let m = pending.pop_front()?;
Some((m.prompt, m.opts))
}
/// One drained user message ready to be appended to history mid-turn.
pub(super) struct DrainedMessage {
pub content: String,
pub metadata: Option<MessageMetadata>,
}
/// Drains the leading run of **non-synthetic** messages, returning them
/// individually (no coalescing). Stops at the first synthetic message, which is
/// left in the queue for the notification path. Used by the running turn to
/// inject newly-queued user input at a round boundary.
pub(super) fn drain_leading_user(
pending: &mut VecDeque<QueuedMessage>,
) -> Vec<DrainedMessage> {
let mut out = Vec::new();
while pending.front().is_some_and(|m| !m.opts.is_synthetic) {
let mut m = pending.pop_front().unwrap();
out.push(DrainedMessage {
content: m.prompt,
metadata: m.opts.metadata.take(),
});
}
out
}
#[cfg(test)]
mod tests {
use super::*;
fn msg(prompt: &str, synthetic: bool) -> QueuedMessage {
QueuedMessage {
prompt: prompt.to_string(),
opts: SendMessageOptions { is_synthetic: synthetic, ..Default::default() },
}
}
#[test]
fn empty_queue_yields_none() {
let mut q = VecDeque::new();
assert!(build_unit(&mut q).is_none());
}
#[test]
fn build_unit_pops_a_single_message() {
let mut q = VecDeque::from(vec![msg("hello", false), msg("also this", false)]);
let (prompt, _) = build_unit(&mut q).unwrap();
assert_eq!(prompt, "hello");
// The second message is left for the round-boundary drain.
assert_eq!(q.len(), 1);
}
#[test]
fn drain_returns_leading_user_messages_individually() {
let mut q = VecDeque::from(vec![msg("a", false), msg("b", false)]);
let drained = drain_leading_user(&mut q);
let contents: Vec<_> = drained.iter().map(|d| d.content.as_str()).collect();
assert_eq!(contents, vec!["a", "b"]);
assert!(q.is_empty());
}
#[test]
fn drain_stops_at_a_synthetic_boundary() {
let mut q = VecDeque::from(vec![
msg("a", false),
msg("b", false),
msg("notification", true),
]);
let drained = drain_leading_user(&mut q);
let contents: Vec<_> = drained.iter().map(|d| d.content.as_str()).collect();
assert_eq!(contents, vec!["a", "b"]);
assert_eq!(q.len(), 1); // the synthetic message is left for the next unit
}
#[test]
fn drain_skips_leading_synthetic() {
let mut q = VecDeque::from(vec![msg("notification", true), msg("user text", false)]);
let drained = drain_leading_user(&mut q);
assert!(drained.is_empty());
assert_eq!(q.len(), 2);
}
fn msg_with_attachment(prompt: &str, path: &str) -> QueuedMessage {
use core_api::message_meta::{Attachment, MessageMetadata};
QueuedMessage {
prompt: prompt.to_string(),
opts: SendMessageOptions {
metadata: Some(MessageMetadata {
attachments: vec![Attachment {
path: path.to_string(),
name: path.to_string(),
mimetype: None,
filesize: None,
}],
..Default::default()
}),
..Default::default()
},
}
}
#[test]
fn drain_preserves_per_message_attachments() {
let mut q = VecDeque::from(vec![
msg_with_attachment("first", "a.pdf"),
msg_with_attachment("second", "b.pdf"),
]);
let drained = drain_leading_user(&mut q);
assert_eq!(drained.len(), 2);
assert_eq!(drained[0].metadata.as_ref().unwrap().attachments[0].path, "a.pdf");
assert_eq!(drained[1].metadata.as_ref().unwrap().attachments[0].path, "b.pdf");
}
}
+843
View File
@@ -0,0 +1,843 @@
use std::collections::HashMap;
use std::sync::atomic::Ordering;
use std::sync::{Arc, OnceLock, Weak};
use std::time::Duration;
use async_trait::async_trait;
use sqlx::SqlitePool;
use tokio::sync::{Mutex, broadcast, mpsc};
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};
mod inbox;
use inbox::{QueuedMessage, SourceInbox, build_unit, drain_leading_user};
use crate::approval::ApprovalManager;
use crate::cron::TaskManager;
use crate::db::{chat_history, chat_llm_tools, chat_sessions, chat_sessions_stack, config, sources};
use crate::events::{GlobalEvent, ServerEvent};
use crate::notification::Notification;
use crate::session::handler::{ChatSessionHandler, InterfaceTool, PendingMsg, PendingUserInput};
use crate::session::manager::ChatSessionManager;
use crate::tools::tool_names as tn;
pub use core_api::chat_hub::{ChatHubApi, ModelCommandOutcome, SendMessageOptions};
pub const HOME_SOURCE_KEY: &str = "source_home";
pub const DEFAULT_HOME_SOURCE: &str = "web";
// Global broadcast channel capacity.
const EVENTS_CAPACITY: usize = 512;
// Central notification queue capacity (inbound from background agents).
const NOTIFY_CAPACITY: usize = 64;
// How long to wait after the first notification before draining, to batch bursts.
const NOTIFY_BATCH_WINDOW_MS: u64 = 200;
// Idle-debounce for per-source message coalescing. 0 = pure coalesce-while-busy
// (a message to an idle source dispatches immediately). Raise it to also batch
// messages sent rapidly to an idle source, at the cost of that latency on the
// first message of a burst.
const SOURCE_COALESCE_DEBOUNCE_MS: u64 = 0;
// ── ChatHub ───────────────────────────────────────────────────────────────────
/// Manages **interactive, user-facing sessions only** (web, mobile, project chats):
/// one live, persistent session per `source`, reachable over WebSocket and addressed
/// by source id through the `sources` table.
///
/// It is **not** a runner for background / non-interactive agents (cron jobs, TIC,
/// sub-agent tasks). Those go through `TaskManager` / `ChatSessionManager` directly and
/// must not be routed here — they are not user-facing, have no broadcast audience, and
/// should not appear in the `sources` table. (Historically this class was misused to
/// drive non-interactive agents; keep that boundary.)
pub struct ChatHub {
db: Arc<SqlitePool>,
session_mgr: Arc<ChatSessionManager>,
pub approval: Arc<ApprovalManager>,
/// Single global broadcast bus. All events from all sources flow here,
/// wrapped in GlobalEvent with source/session_id tags. Subscribers filter.
global_tx: broadcast::Sender<GlobalEvent>,
/// Central inbound notification queue from background agents.
/// Consumer task is spawned in new() and drains this channel.
notify_tx: mpsc::Sender<Notification>,
/// TaskManager reference for injecting execute_task into interactive sessions.
/// Set via set_task_mgr() after construction (breaks circular dep with cron).
task_mgr: std::sync::OnceLock<Arc<TaskManager>>,
/// Per-source input inboxes (coalescing + FIFO ordering). Created lazily on the
/// first message for a source; each spawns one consumer task.
inboxes: Mutex<HashMap<String, Arc<SourceInbox>>>,
/// Weak self-reference, set in `new()`, so lazily-spawned source consumers can
/// reach back into the hub to dispatch turns.
me: OnceLock<Weak<Self>>,
/// Shutdown token, used to stop lazily-spawned source consumers.
shutdown: CancellationToken,
/// Per-source pinned LLM client (e.g. set via `/model` or the web dropdown).
/// Keyed by source id; value is a `client_names()` entry (`"auto"` or a
/// model name). When absent the caller AUTO-resolves. In-memory only: a
/// server restart clears all pins (intentional for the MVP).
selected_clients: Mutex<HashMap<String, String>>,
}
impl ChatHub {
pub fn new(
db: Arc<SqlitePool>,
session_mgr: Arc<ChatSessionManager>,
approval: Arc<ApprovalManager>,
global_tx: broadcast::Sender<GlobalEvent>,
shutdown: CancellationToken,
) -> Arc<Self> {
let (notify_tx, notify_rx) = mpsc::channel::<Notification>(NOTIFY_CAPACITY);
let hub = Arc::new(Self {
db,
session_mgr,
approval,
global_tx,
notify_tx,
task_mgr: std::sync::OnceLock::new(),
inboxes: Mutex::new(HashMap::new()),
me: OnceLock::new(),
shutdown: shutdown.clone(),
selected_clients: Mutex::new(HashMap::new()),
});
// Store a weak self-reference for lazily-spawned source consumers.
let _ = hub.me.set(Arc::downgrade(&hub));
// Spawn the background consumer with a Weak reference so it doesn't
// prevent ChatHub from being dropped on shutdown.
tokio::spawn(Self::notification_consumer(Arc::downgrade(&hub), notify_rx, shutdown));
hub
}
/// Called once after TaskManager is built (breaks circular dep: TaskManager needs
/// ChatSessionManager, ChatHub needs TaskManager for execute_task injection).
pub fn set_task_mgr(&self, task_mgr: Arc<TaskManager>) {
let _ = self.task_mgr.set(task_mgr);
}
// ── Public API ────────────────────────────────────────────────────────────
/// Register a source. No-op for duplicate registrations.
/// With the global bus, registration no longer creates a per-source channel.
pub async fn register(&self, source_id: &str) {
info!(source_id, "ChatHub: source registered");
}
/// Enqueue a user message for a source. Returns immediately once queued; the
/// turn runs asynchronously on the source's consumer task, which coalesces
/// messages that pile up during an in-flight turn into a single follow-up turn
/// (see `inbox`). Creates the source's inbox (and consumer) lazily on first use.
/// Turn errors surface via the `Error` event on the broadcast bus, not this
/// return value.
pub async fn send_message(
&self,
source_id: &str,
prompt: &str,
opts: SendMessageOptions,
) -> anyhow::Result<()> {
let inbox = self.get_or_spawn_inbox(source_id).await;
inbox.pending.lock().await.push_back(QueuedMessage {
prompt: prompt.to_string(),
opts,
});
inbox.notify.notify_one();
Ok(())
}
/// Returns the source's inbox, creating it (and spawning its consumer) on first use.
async fn get_or_spawn_inbox(&self, source_id: &str) -> Arc<SourceInbox> {
let mut inboxes = self.inboxes.lock().await;
if let Some(inbox) = inboxes.get(source_id) {
return Arc::clone(inbox);
}
let inbox = Arc::new(SourceInbox::default());
inboxes.insert(source_id.to_string(), Arc::clone(&inbox));
let weak = self.me.get().expect("ChatHub::me must be set in new()").clone();
tokio::spawn(Self::source_consumer(
weak,
source_id.to_string(),
Arc::clone(&inbox),
self.shutdown.clone(),
));
info!(source_id, "ChatHub: source inbox + consumer spawned");
inbox
}
/// Runs one LLM turn for a coalesced unit: resolves session/handler, bridges
/// events to the global bus, injects `execute_task`, and calls `handle_message`
/// (which takes the per-session `processing` lock).
async fn dispatch_turn(
&self,
source_id: &str,
prompt: &str,
opts: SendMessageOptions,
// Live user-input source for this turn (the source's inbox). The running
// turn drains it at each round boundary to inject messages queued while it
// was busy. `None` for synthetic turns, which never inject.
pending_input: Option<Arc<dyn PendingUserInput>>,
) -> anyhow::Result<()> {
let agent_id = opts.agent_id.as_deref().unwrap_or("main");
let session_id = self.get_or_create_session(source_id, agent_id).await?;
let source_tag = source_id.to_string();
// Bridge mpsc from handle_message → global broadcast, tagging with source/session.
let tx = Self::bridge_to_global(self.global_tx.clone(), source_tag, session_id);
// get_or_create_handler is idempotent; we call it early to read the
// session's RunContext so it can be inherited by any task spawned here.
let handler = self.session_mgr.get_or_create_handler(session_id).await?;
let run_context_json = handler.run_context_json().await;
// Inject execute_task as an InterfaceTool for all interactive sessions.
// session_id and run_context_json are captured so tasks inherit the parent context.
let mut interface_tools = opts.interface_tools;
if let Some(task_mgr) = self.task_mgr.get() {
interface_tools.push(
crate::tools::cron_jobs::build_execute_task_interface_tool(
Arc::clone(task_mgr),
session_id,
run_context_json,
)
);
}
handler.handle_message(
prompt,
opts.client_name,
opts.extra_system_context,
opts.extra_system_dynamic,
opts.tail_reminder,
interface_tools,
opts.system_substitutions,
tx,
opts.is_synthetic,
opts.metadata,
pending_input,
).await
}
/// Returns the session handler for the source's active session, creating one lazily if needed.
pub async fn session_handler(&self, source_id: &str) -> anyhow::Result<Arc<ChatSessionHandler>> {
let session_id = self.get_or_create_session(source_id, "main").await?;
self.session_mgr.get_or_create_handler(session_id).await
}
/// Returns the handler for a specific `session_id`, creating one lazily if needed.
/// Used to resolve a pending tool against the session that actually owns it,
/// independent of any source's "active" session.
pub async fn handler_for_session(&self, session_id: i64) -> anyhow::Result<Arc<ChatSessionHandler>> {
self.session_mgr.get_or_create_handler(session_id).await
}
/// Ensures a persistent, interactive session exists for `source`, created with
/// `agent_id` and the given `run_context`.
///
/// If a session already exists for the source it is returned as-is, unless `reset`
/// is set — in which case the existing session is discarded and a fresh one is
/// created (and a `NewSession` event is broadcast so connected clients reset).
///
/// This is the single entry point for the source→session mapping ChatHub owns.
/// Note: `agent_id`/`run_context` only take effect when a session is actually
/// created; on reuse the existing session keeps its original agent and context.
pub async fn provision_session(
&self,
source_id: &str,
agent_id: &str,
run_context: Option<&crate::run_context::RunContext>,
reset: bool,
) -> anyhow::Result<i64> {
// A reset discards the current session; drop any messages queued for it.
if reset {
self.clear_inbox(source_id).await;
}
if !reset {
if let Some(sid) = sources::active_session_id(&self.db, source_id).await? {
return Ok(sid);
}
}
let (session_id, _) = self.session_mgr
.create_session(agent_id, source_id, true, false, run_context)
.await?;
sources::upsert(&self.db, source_id, session_id).await?;
info!(source_id, session_id, agent_id, reset, "ChatHub: session provisioned");
if reset {
let _ = self.global_tx.send(GlobalEvent {
source: Some(source_id.to_string()),
session_id: Some(session_id),
event: ServerEvent::NewSession { session_id },
});
}
Ok(session_id)
}
/// Create a new session for the source, discarding the previous one.
/// Thin wrapper over `provision_session` preserving the default `main` agent
/// (kept for the `ChatHubApi` trait and generic callers).
pub async fn clear(&self, source_id: &str) -> anyhow::Result<i64> {
self.provision_session(source_id, "main", None, true).await
}
/// Subscribe to the global event bus. The `source_id` parameter is accepted
/// for API compatibility but filtering by source is the caller's responsibility.
pub fn events(&self, _source_id: &str) -> broadcast::Receiver<GlobalEvent> {
self.global_tx.subscribe()
}
/// Emit an event directly on the global bus (for system events without a session).
pub fn emit(&self, event: GlobalEvent) {
let _ = self.global_tx.send(event);
}
/// Set which source is the "home" for background agent notifications.
pub async fn set_home(&self, source_id: &str) -> anyhow::Result<()> {
config::set(&self.db, HOME_SOURCE_KEY, source_id).await?;
info!(source_id, "ChatHub: home source set");
Ok(())
}
/// Returns the current home source id, falling back to `web` if not configured.
pub async fn home_source(&self) -> anyhow::Result<String> {
Ok(config::get(&self.db, HOME_SOURCE_KEY)
.await?
.unwrap_or_else(|| DEFAULT_HOME_SOURCE.to_string()))
}
/// Returns token usage for the last message in the source's active session.
/// Returns `(input_tokens, output_tokens)` — both are `None` when no
/// messages exist or the provider did not report usage.
pub async fn context_info(&self, source_id: &str) -> anyhow::Result<(Option<i64>, Option<i64>)> {
let session_id = self.get_or_create_session(source_id, "main").await?;
let stack = match chat_sessions_stack::active_for_session(&self.db, session_id).await? {
Some(s) => s,
None => return Ok((None, None)),
};
let last = chat_history::last_message_for_stack(&self.db, stack.id).await?;
Ok(last.map_or((None, None), |m| (m.input_tokens, m.output_tokens)))
}
/// Total spend (USD) of the source's active session, including synchronous
/// sub-agent frames and excluding asynchronous tasks (which run in their own
/// session). `None` when no provider reported a cost.
pub async fn cost_info(&self, source_id: &str) -> anyhow::Result<Option<f64>> {
let session_id = self.get_or_create_session(source_id, "main").await?;
chat_history::total_cost_for_session(&self.db, session_id).await
}
/// Force compaction of the source's active session history.
/// Bypasses the token threshold; returns `true` if compaction occurred.
pub async fn force_compact(&self, source_id: &str) -> anyhow::Result<bool> {
let handler = self.session_handler(source_id).await?;
handler.force_compact().await
}
/// Resume any interrupted turn for a source's active session.
/// Calls `resume_turn` which re-executes pending tool calls (approval or
/// clarification) and re-runs the LLM loop if needed.
/// Safe to call unconditionally — returns immediately if there is nothing to resume.
/// Events are published to the global broadcast bus so existing subscribers
/// (e.g. Telegram's persistent_forwarder) receive them without a WS connection.
pub async fn resume(&self, source_id: &str) -> anyhow::Result<()> {
let session_id = match sources::active_session_id(&self.db, source_id).await? {
Some(sid) => sid,
None => return Ok(()), // no prior session, nothing to resume
};
self.resume_session(session_id).await
}
/// Resume an interrupted turn for a specific `session_id` (post-restart recovery
/// or after a manual approval resolve), independent of any source's active session.
/// Injects `execute_task` so a pending sub-agent task can be re-dispatched, and
/// bridges events to the global bus so the reconnected client still sees them.
pub async fn resume_session(&self, session_id: i64) -> anyhow::Result<()> {
// Source tag drives per-source event filtering for connected clients.
let source = chat_sessions::find_by_id(&self.db, session_id).await?
.map(|s| s.source)
.unwrap_or_else(|| "web".to_string());
let tx = Self::bridge_to_global(self.global_tx.clone(), source, session_id);
let handler = self.session_mgr.get_or_create_handler(session_id).await?;
let interface_tools = self.execute_task_tools(session_id, &handler).await;
handler.resume_turn(None, None, interface_tools, tx).await
}
/// Builds the `execute_task` interface tool for a session, mirroring the injection
/// done for live turns (`run_agent_turn`). Empty when no TaskManager is configured
/// so `execute_task mode=async` can be rebuilt by `build_execution` during resume.
async fn execute_task_tools(
&self,
session_id: i64,
handler: &Arc<ChatSessionHandler>,
) -> Vec<InterfaceTool> {
let mut tools = Vec::new();
if let Some(task_mgr) = self.task_mgr.get() {
let run_context_json = handler.run_context_json().await;
tools.push(crate::tools::cron_jobs::build_execute_task_interface_tool(
Arc::clone(task_mgr),
session_id,
run_context_json,
));
}
tools
}
/// Queue a structured notification from a background agent.
/// The consumer task aggregates pending notifications and dispatches them to the home source.
pub async fn notify(&self, note: Notification) -> anyhow::Result<()> {
if self.notify_tx.send(note).await.is_err() {
warn!("ChatHub::notify: notification queue full or receiver dropped");
}
Ok(())
}
/// Synchronous variant of `notify` for use inside `Tool::execute` (sync context).
/// Uses `try_send` — drops the notification if the channel is full rather than blocking.
pub fn notify_sync(&self, note: Notification) {
if self.notify_tx.try_send(note).is_err() {
warn!("ChatHub::notify_sync: notification channel full or closed — notification dropped");
}
}
/// Revoke all session-scoped MCP grants for a source's active session.
/// The next LLM turn will start with no MCP servers activated.
pub async fn reset_mcp(&self, source_id: &str) -> anyhow::Result<()> {
let session_id = self.get_or_create_session(source_id, "main").await?;
crate::db::session_mcp_grants::revoke_all(&self.db, session_id).await?;
info!(source_id, session_id, "ChatHub: MCP grants reset");
Ok(())
}
// ── Per-source pinned LLM client ─────────────────────────────────────────
//
// Backend-owned state: every UI mutation (Telegram `/model`, web `/model`,
// web dropdown change) funnels through `set_selected_client`, which then
// broadcasts `ClientSelected` to all clients of the source. The web dropdown
// and mobile select read this event to stay in sync — the backend is the
// single source of truth. Pattern is intentionally generic so future
// per-source toggles (e.g. reasoning level) can mirror it.
/// Returns `(models, default)` — `models` is the ordered list of usable
/// client names (`"auto"` first, then models by priority/name), `default`
/// is the configured default client name.
pub async fn list_clients(&self) -> (Vec<String>, String) {
let mgr = self.session_mgr.llm_manager();
(mgr.client_names().await, mgr.default_name().await)
}
/// Returns the client name pinned for the source, or `None` when unset
/// (the caller should fall back to AUTO resolution).
pub async fn get_selected_client(&self, source_id: &str) -> Option<String> {
self.selected_clients.lock().await.get(source_id).cloned()
}
/// Pin a client name for the source and broadcast `ClientSelected`.
/// `client` should be a `list_clients()` entry (`"auto"` or a model name).
pub async fn set_selected_client(&self, source_id: &str, client: String) {
info!(source_id, client = %client, "ChatHub: selected client set");
self.selected_clients.lock().await.insert(source_id.to_string(), client.clone());
self.emit(GlobalEvent {
source: Some(source_id.to_string()),
session_id: None,
event: ServerEvent::ClientSelected { client },
});
}
/// Clear any pinned client for the source (revert to AUTO) and broadcast
/// `ClientSelected { client: "auto" }`.
pub async fn clear_selected_client(&self, source_id: &str) {
info!(source_id, "ChatHub: selected client cleared (auto)");
self.selected_clients.lock().await.remove(source_id);
self.emit(GlobalEvent {
source: Some(source_id.to_string()),
session_id: None,
event: ServerEvent::ClientSelected { client: "auto".to_string() },
});
}
/// Snapshot of the model list with the per-source current selection marked.
/// Returns `(index, name, is_current)` tuples so call sites can render
/// HTML (Telegram) or Markdown (web) without re-querying the LLM manager
/// or the pin store.
pub async fn list_clients_marked(&self, source_id: &str) -> Vec<(usize, String, bool)> {
let (models, _default) = self.list_clients().await;
let current = self.get_selected_client(source_id).await
.unwrap_or_else(|| "auto".to_string());
models.into_iter()
.enumerate()
.map(|(i, name)| (i, name.clone(), name == current))
.collect()
}
/// Apply a `/model {arg}` command: resolve the argument, mutate the
/// per-source pinned client (broadcasting `ClientSelected`), return a
/// structured outcome. Business logic is centralised here so Telegram and
/// web share a single code path; only the formatting differs.
pub async fn apply_model_command(
&self,
source_id: &str,
arg: &str,
) -> ModelCommandOutcome {
let (models, _default) = self.list_clients().await;
match core_api::chat_hub::resolve_list_arg(&models, arg) {
Ok(Some(client)) => {
let name = client.clone();
self.set_selected_client(source_id, client).await;
ModelCommandOutcome::Set(name)
}
Ok(None) => {
self.clear_selected_client(source_id).await;
ModelCommandOutcome::Cleared
}
Err(msg) => ModelCommandOutcome::Error(msg),
}
}
/// Cancel the active LLM turn for the source's session, clearing any pending
/// approvals and clarification questions. No-op if no session is active.
pub async fn cancel(&self, source_id: &str) {
// Drop queued-but-not-yet-dispatched messages so /stop clears the backlog
// too, not just the in-flight turn.
self.clear_inbox(source_id).await;
match self.session_handler(source_id).await {
Ok(handler) => {
handler.cancel();
handler.cancel_pending_approvals().await;
handler.cancel_pending_questions().await;
info!(source_id, "ChatHub: cancel requested");
}
Err(e) => {
warn!(source_id, error = %e, "ChatHub::cancel: no session to cancel");
}
}
}
/// Approve a pending tool-call approval request.
pub async fn approve(&self, request_id: i64) {
self.approval.approve(request_id).await;
}
/// Reject a pending tool-call approval request.
pub async fn reject(&self, request_id: i64, note: String) {
self.approval.reject(request_id, note).await;
}
// ── Private helpers ───────────────────────────────────────────────────────
/// Spawn a bridge task that forwards events from an mpsc channel to the
/// global broadcast bus, tagging each event with `source` and `session_id`.
fn bridge_to_global(
global_tx: broadcast::Sender<GlobalEvent>,
source: String,
session_id: i64,
) -> mpsc::Sender<ServerEvent> {
let (tx, mut rx) = mpsc::channel::<ServerEvent>(EVENTS_CAPACITY);
tokio::spawn(async move {
tracing::debug!(%source, session_id, "ChatHub: bridge task started");
while let Some(event) = rx.recv().await {
tracing::debug!(%source, session_id, event_type = event.type_name(), "ChatHub: bridge forwarding event");
let _ = global_tx.send(GlobalEvent {
source: Some(source.clone()),
session_id: Some(session_id),
event,
});
}
tracing::debug!(%source, session_id, "ChatHub: bridge task ended");
});
tx
}
async fn get_or_create_session(&self, source_id: &str, agent_id: &str) -> anyhow::Result<i64> {
if let Some(sid) = sources::active_session_id(&self.db, source_id).await? {
return Ok(sid);
}
let (session_id, _) = self.session_mgr.create_session(agent_id, source_id, true, false, None).await?;
sources::upsert(&self.db, source_id, session_id).await?;
info!(source_id, session_id, "ChatHub: session created lazily");
Ok(session_id)
}
// ── Per-source inbox consumer ─────────────────────────────────────────────
/// Per-source consumer: drains and coalesces queued messages, running one turn
/// at a time. Spawned lazily by `get_or_spawn_inbox`; lives until shutdown.
async fn source_consumer(
hub: Weak<Self>,
source_id: String,
inbox: Arc<SourceInbox>,
shutdown: CancellationToken,
) {
info!(%source_id, "ChatHub: source consumer started");
loop {
tokio::select! {
_ = shutdown.cancelled() => break,
_ = inbox.notify.notified() => {}
}
// Optional idle-batching window (0 = disabled).
if SOURCE_COALESCE_DEBOUNCE_MS > 0 {
tokio::time::sleep(Duration::from_millis(SOURCE_COALESCE_DEBOUNCE_MS)).await;
}
// Pop one message to seed a turn, then dispatch it. Messages that arrive
// while the turn runs are injected live at its round boundaries (the turn
// drains `pending` itself via the `PendingUserInput` handle below); only
// messages that arrive after the turn's last boundary remain here and seed
// the next turn on a following iteration.
loop {
let (unit, epoch) = {
let mut pending = inbox.pending.lock().await;
let epoch = inbox.cancel_epoch.load(Ordering::Acquire);
(build_unit(&mut pending), epoch)
};
let Some((prompt, opts)) = unit else { break };
let Some(hub) = hub.upgrade() else { return };
// A /stop between draining and dispatching bumps cancel_epoch and
// clears pending — drop this now-stale unit.
if inbox.cancel_epoch.load(Ordering::Acquire) != epoch {
continue;
}
// Live-injection source for this turn (real user turns only).
let pending_input: Option<Arc<dyn PendingUserInput>> = (!opts.is_synthetic)
.then(|| Arc::new(InboxUserInput(Arc::clone(&inbox))) as Arc<dyn PendingUserInput>);
// Run the turn on a dedicated task and await its handle. Isolating it
// means a panic inside the turn (e.g. a UTF-8 boundary slice on a
// tool-result preview) surfaces here as a JoinError and is logged,
// instead of unwinding the consumer task and silently killing this
// source's chat (new messages would then enqueue and never dispatch).
let hub_turn = Arc::clone(&hub);
let src = source_id.clone();
let turn = tokio::spawn(async move {
hub_turn.dispatch_turn(&src, &prompt, opts, pending_input).await
});
match turn.await {
Ok(Ok(())) => {}
Ok(Err(e)) => error!(%source_id, error = %e, "ChatHub: source turn failed"),
Err(e) => error!(%source_id, error = %e, "ChatHub: source turn panicked — consumer surviving"),
}
}
}
info!(%source_id, "ChatHub: source consumer stopped");
}
/// Clears a source's pending queue and bumps its cancel epoch (so a unit the
/// consumer drained just before a `/stop` is dropped instead of dispatched).
/// No-op if the source has no inbox yet.
async fn clear_inbox(&self, source_id: &str) {
if let Some(inbox) = self.inboxes.lock().await.get(source_id) {
inbox.pending.lock().await.clear();
inbox.cancel_epoch.fetch_add(1, Ordering::Release);
}
}
// ── Notification consumer ─────────────────────────────────────────────────
/// Background task: drains the central notification queue and dispatches
/// aggregated briefings to the home source as synthetic user messages.
///
/// Serialisation with active LLM turns is free: `ChatSessionHandler::handle_message`
/// holds `processing: Mutex<()>` for the duration of a turn, so `send_message`
/// below blocks naturally until the turn completes.
async fn notification_consumer(hub: Weak<Self>, mut rx: mpsc::Receiver<Notification>, shutdown: CancellationToken) {
info!("ChatHub: notification consumer started");
loop {
// Block until at least one notification arrives (or shutdown signal).
let first = tokio::select! {
_ = shutdown.cancelled() => {
info!("ChatHub: notification consumer shutdown");
break;
}
msg = rx.recv() => match msg {
Some(n) => n,
None => break, // notify_tx dropped — ChatHub is shutting down
}
};
// Brief window to let burst notifications accumulate before dispatching.
tokio::time::sleep(Duration::from_millis(NOTIFY_BATCH_WINDOW_MS)).await;
// Drain everything else that arrived during the window.
let mut notes = vec![first];
while let Ok(n) = rx.try_recv() {
notes.push(n);
}
let hub = match hub.upgrade() {
Some(h) => h,
None => break, // ChatHub dropped
};
let home = match hub.home_source().await {
Ok(h) => h,
Err(e) => { error!(error = %e, "notification consumer: home_source failed"); continue; }
};
let count = notes.len();
// Build a synthetic assistant message with a reasoning trace and a
// pre-completed read_notification tool call carrying the notifications as results.
// The agent is then woken via resume() — resume_turn sees the tool calls on
// the last assistant message and runs the LLM loop so the agent can respond.
let result_json = serde_json::to_string(&notes).unwrap_or_else(|_| "[]".to_string());
let session_id = match hub.get_or_create_session(&home, "main").await {
Ok(sid) => sid,
Err(e) => { error!(error = %e, "notification consumer: get_or_create_session failed"); continue; }
};
let stack = match chat_sessions_stack::active_for_session(&hub.db, session_id).await {
Ok(Some(s)) => s,
Ok(None) => { error!(session_id, "notification consumer: no active stack"); continue; }
Err(e) => { error!(error = %e, "notification consumer: active_for_session failed"); continue; }
};
let assistant_id = match chat_history::append(
&hub.db, stack.id, &chat_history::Role::Assistant,
"", true,
Some("The system signaled pending notifications. Let me read them and surface anything relevant to the user."),
).await {
Ok(id) => id,
Err(e) => { error!(error = %e, "notification consumer: append assistant failed"); continue; }
};
let tool_call_id = match chat_llm_tools::append(
&hub.db, assistant_id, tn::READ_NOTIFICATION, "{}",
).await {
Ok(id) => id,
Err(e) => { error!(error = %e, "notification consumer: append tool call failed"); continue; }
};
if let Err(e) = chat_llm_tools::complete(&hub.db, tool_call_id, &result_json, "json").await {
error!(error = %e, "notification consumer: complete tool call failed"); continue;
}
info!(home_source = %home, count, "ChatHub: dispatching notifications via read_notification");
if let Err(e) = hub.resume(&home).await {
error!(error = %e, "notification consumer: resume failed");
}
}
info!("ChatHub: notification consumer stopped");
}
}
// ── Live user-input source ──────────────────────────────────────────────────
/// Adapts a source's `SourceInbox` to the handler's `PendingUserInput` trait so a
/// running turn can drain newly-queued user messages at its round boundaries.
struct InboxUserInput(Arc<SourceInbox>);
#[async_trait]
impl PendingUserInput for InboxUserInput {
async fn drain_user(&self) -> Vec<PendingMsg> {
let mut pending = self.0.pending.lock().await;
drain_leading_user(&mut pending)
.into_iter()
.map(|d| PendingMsg { content: d.content, metadata: d.metadata })
.collect()
}
}
// ── ChatHubApi impl ───────────────────────────────────────────────────────────
#[async_trait]
impl ChatHubApi for ChatHub {
async fn register(&self, source_id: &str) {
self.register(source_id).await
}
async fn send_message(
&self,
source_id: &str,
prompt: &str,
opts: SendMessageOptions,
) -> anyhow::Result<()> {
self.send_message(source_id, prompt, opts).await
}
async fn clear(&self, source_id: &str) -> anyhow::Result<i64> {
self.clear(source_id).await
}
fn events(&self, source_id: &str) -> broadcast::Receiver<GlobalEvent> {
self.events(source_id)
}
async fn set_home(&self, source_id: &str) -> anyhow::Result<()> {
self.set_home(source_id).await
}
async fn context_info(&self, source_id: &str) -> anyhow::Result<(Option<i64>, Option<i64>)> {
self.context_info(source_id).await
}
async fn cost_info(&self, source_id: &str) -> anyhow::Result<Option<f64>> {
self.cost_info(source_id).await
}
async fn force_compact(&self, source_id: &str) -> anyhow::Result<bool> {
self.force_compact(source_id).await
}
async fn resume(&self, source_id: &str) -> anyhow::Result<()> {
self.resume(source_id).await
}
async fn approve(&self, request_id: i64) {
self.approve(request_id).await
}
async fn reject(&self, request_id: i64, note: String) {
self.reject(request_id, note).await
}
async fn resolve_question(&self, source_id: &str, request_id: i64, answer: String) {
if let Ok(handler) = self.session_handler(source_id).await {
handler.resolve_question(request_id, answer).await;
} else {
warn!(source_id, request_id, "ChatHubApi::resolve_question: no session handler");
}
}
async fn cancel(&self, source_id: &str) {
self.cancel(source_id).await
}
async fn reset_mcp(&self, source_id: &str) -> anyhow::Result<()> {
self.reset_mcp(source_id).await
}
async fn list_clients(&self) -> (Vec<String>, String) {
self.list_clients().await
}
async fn get_selected_client(&self, source_id: &str) -> Option<String> {
self.get_selected_client(source_id).await
}
async fn set_selected_client(&self, source_id: &str, client: String) {
self.set_selected_client(source_id, client).await;
}
async fn clear_selected_client(&self, source_id: &str) {
self.clear_selected_client(source_id).await;
}
async fn list_clients_marked(
&self,
source_id: &str,
) -> Vec<(usize, String, bool)> {
self.list_clients_marked(source_id).await
}
async fn apply_model_command(
&self,
source_id: &str,
arg: &str,
) -> ModelCommandOutcome {
self.apply_model_command(source_id, arg).await
}
}
+162
View File
@@ -0,0 +1,162 @@
//! Transparent logging wrapper for any [`ChatbotClient`].
//!
//! [`LoggingChatbotClient`] intercepts every `chat_with_tools` call, captures
//! the raw HTTP request/response from the inner provider via `chat_with_tools_raw`,
//! then persists a row to `llm_requests` asynchronously (fire-and-forget).
//!
//! The LLM loop is completely unaware of this: it only holds an
//! `Arc<dyn ChatbotClient>` and calls `chat_with_tools` as usual.
use std::sync::Arc;
use std::time::Instant;
use async_trait::async_trait;
use serde_json::Value;
use sqlx::SqlitePool;
use tracing::warn;
use crate::db::llm_requests;
use super::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message};
// ─────────────────────────────────────────────────────────────────────────────
/// Controls which parts of the HTTP exchange are persisted per row.
#[derive(Debug, Clone, Copy)]
pub struct LogSaveFlags {
pub request_payload: bool,
pub response_payload: bool,
pub request_headers: bool,
pub response_headers: bool,
}
impl Default for LogSaveFlags {
fn default() -> Self {
Self { request_payload: true, response_payload: true, request_headers: true, response_headers: true }
}
}
pub struct LoggingChatbotClient {
inner: Arc<dyn ChatbotClient>,
pool: Arc<SqlitePool>,
model_name: String,
flags: LogSaveFlags,
}
impl LoggingChatbotClient {
pub fn new(
inner: Arc<dyn ChatbotClient>,
pool: Arc<SqlitePool>,
model_name: impl Into<String>,
flags: LogSaveFlags,
) -> Self {
Self { inner, pool, model_name: model_name.into(), flags }
}
}
#[async_trait]
impl ChatbotClient for LoggingChatbotClient {
/// Passthrough — logging only applies to the tool-calling path.
async fn chat(
&self,
messages: &[Message],
options: &ChatOptions,
) -> anyhow::Result<ChatResponse> {
self.inner.chat(messages, options).await
}
/// Intercepts the call, delegates to `inner.chat_with_tools_raw` to capture
/// HTTP wire data, then spawns a fire-and-forget DB write before returning.
async fn chat_with_tools(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<LlmTurn> {
let start = Instant::now();
let result = self.inner.chat_with_tools_raw(messages, tools, options).await;
let duration_ms = start.elapsed().as_millis() as i64;
let session_id = options.session_id;
let stack_id = options.stack_id;
let model_name = self.model_name.clone();
let pool = Arc::clone(&self.pool);
match result {
Ok((turn, meta)) => {
let (input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens) = match &turn {
LlmTurn::Message(r) => (r.input_tokens, r.output_tokens, r.cache_read_tokens, r.cache_creation_tokens),
LlmTurn::ToolCalls { input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens, .. } =>
(*input_tokens, *output_tokens, *cache_read_tokens, *cache_creation_tokens),
};
let meta = meta.unwrap_or_default();
let flags = self.flags;
let request_json = if flags.request_payload {
meta.request_body.map(|v| v.to_string()).unwrap_or_default()
} else { String::new() };
let request_headers = if flags.request_headers { meta.request_headers.map(|v| v.to_string()) } else { None };
let response_json = if flags.response_payload { meta.response_body.map(|v| v.to_string()) } else { None };
let response_headers = if flags.response_headers { meta.response_headers.map(|v| v.to_string()) } else { None };
tokio::spawn(async move {
if let Err(e) = llm_requests::insert(&pool, llm_requests::LlmRequestRow {
session_id,
stack_id,
model_name,
request_json,
request_headers,
response_json,
response_headers,
error_text: None,
input_tokens: input_tokens.map(|n| n as i64),
output_tokens: output_tokens.map(|n| n as i64),
duration_ms,
cache_read_tokens: cache_read_tokens.map(|n| n as i64),
cache_creation_tokens: cache_creation_tokens.map(|n| n as i64),
}).await {
warn!(error = %e, "llm_requests: failed to insert log row");
}
});
Ok(turn)
}
Err(e) => {
let error_text = e.to_string();
tokio::spawn(async move {
if let Err(log_err) = llm_requests::insert(&pool, llm_requests::LlmRequestRow {
session_id,
stack_id,
model_name,
request_json: String::new(),
request_headers: None,
response_json: None,
response_headers: None,
error_text: Some(error_text),
input_tokens: None,
output_tokens: None,
duration_ms,
cache_read_tokens: None,
cache_creation_tokens: None,
}).await {
warn!(error = %log_err, "llm_requests: failed to insert error log row");
}
});
Err(e)
}
}
}
/// Expose raw metadata so this wrapper can itself be wrapped if needed.
async fn chat_with_tools_raw(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
self.inner.chat_with_tools_raw(messages, tools, options).await
}
}
+7
View File
@@ -0,0 +1,7 @@
pub mod logging;
// Re-export from the independent llm-client crate.
pub use llm_client::{
ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Message, ToolCall,
anthropic, lm_studio, ollama, openai,
};
+107
View File
@@ -0,0 +1,107 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicI64, Ordering};
use chrono::Utc;
use serde::Serialize;
use tokio::sync::{broadcast, oneshot};
use tracing::info;
use crate::events::{GlobalEvent, ServerEvent};
use crate::pending_registry::PendingRegistry;
#[derive(Debug, Clone, Serialize)]
pub struct PendingClarificationInfo {
pub request_id: i64,
pub session_id: i64,
pub agent_id: String,
pub source: String,
pub context_label: Option<String>,
pub title: String,
pub question: String,
pub suggested_answers: Vec<String>,
pub created_at: String,
}
pub struct ClarificationManager {
/// Shared pending-request plumbing (map + oneshot). Keyed by `request_id`.
registry: PendingRegistry<PendingClarificationInfo, String>,
next_id: AtomicI64,
/// Global event bus sender, mirroring `ApprovalManager`. Used to broadcast
/// `ClarificationRequested` / `ClarificationResolved` so Inbox subscribers
/// (e.g. the mobile-connector plugin) can re-snapshot.
event_tx: broadcast::Sender<GlobalEvent>,
}
impl ClarificationManager {
pub fn new(event_tx: broadcast::Sender<GlobalEvent>) -> Arc<Self> {
Arc::new(Self {
registry: PendingRegistry::new(),
next_id: AtomicI64::new(1),
event_tx,
})
}
pub async fn register(
&self,
session_id: i64,
agent_id: &str,
source: &str,
context_label: Option<&str>,
title: &str,
question: &str,
suggested_answers: Vec<String>,
) -> (i64, oneshot::Receiver<String>) {
let request_id = self.next_id.fetch_add(1, Ordering::SeqCst);
let info = PendingClarificationInfo {
request_id,
session_id,
agent_id: agent_id.to_string(),
source: source.to_string(),
context_label: context_label.map(str::to_string),
title: title.to_string(),
question: question.to_string(),
suggested_answers,
created_at: Utc::now().to_rfc3339(),
};
let rx = self.registry.insert(request_id, info).await;
info!(session_id, agent = agent_id, source, request_id, "clarification: pending registered");
// Broadcast on the global bus; counterpart of the per-session
// `AgentQuestion` WS event.
let _ = self.event_tx.send(GlobalEvent {
source: Some(source.to_string()),
session_id: Some(session_id),
event: ServerEvent::ClarificationRequested {
request_id,
title: title.to_string(),
},
});
(request_id, rx)
}
pub async fn resolve(&self, request_id: i64, answer: String) -> bool {
match self.registry.resolve(request_id, answer).await {
Some(info) => {
info!(request_id, "clarification: resolved");
let _ = self.event_tx.send(GlobalEvent {
source: Some(info.source),
session_id: Some(info.session_id),
event: ServerEvent::ClarificationResolved { request_id },
});
true
}
None => false,
}
}
pub async fn list_pending(&self) -> Vec<PendingClarificationInfo> {
let mut items = self.registry.list().await;
items.sort_by(|a, b| a.created_at.cmp(&b.created_at));
items
}
pub async fn cancel_for_session(&self, session_id: i64) {
self.registry.remove_where(|i| i.session_id == session_id).await;
}
}
+115
View File
@@ -0,0 +1,115 @@
//! File-based custom slash commands.
//!
//! Each command lives in `commands/<name>/` with a `meta.json` manifest and a
//! `COMMAND.md` template — mirroring the `agents/<id>/` layout. Files are read at
//! request time so edits take effect without a restart (like `agents::load_prompt`).
//!
//! A recognised `/command` expands its `COMMAND.md` template (interpolating the
//! user's arguments into `{{args}}` / `{{prompt}}`) into a **normal user message on
//! the `main` session**, so the turn stays fully interactive: the model can ask
//! questions, iterate, and dispatch sub-agents exactly as in any other turn.
//!
//! [`CommandApi`] (in `core-api`) is the capability trait plugins depend on; this
//! manager is its only implementation.
use core_api::command::{CommandApi, CommandInfo, ResolvedCommand, expand_template};
use serde::Deserialize;
use tracing::warn;
const COMMANDS_DIR: &str = "commands";
/// Command names that collide with the hard-coded system commands in the WS handler.
/// System commands are matched first, so a same-named custom command is unreachable;
/// these are also filtered out of discovery/listing to avoid dead entries in `/help`
/// and the autocomplete.
const RESERVED: &[&str] = &[
"clear", "new", "help", "context", "cost", "compact",
"resettools", "models", "model", "sethome", "stop",
];
/// The `meta.json` manifest of a custom command.
#[derive(Debug, Clone, Deserialize)]
struct RawMeta {
description: String,
#[serde(default = "default_true")]
enabled: bool,
}
fn default_true() -> bool { true }
/// File-based manager for custom slash commands. Owned by `Skald` (wrapped in
/// `Arc` so it can be shared with plugins as `Arc<dyn CommandApi>`); stateless — it
/// re-reads `commands/` on each call, so edits take effect without a restart.
pub struct LlmCommandManager;
impl LlmCommandManager {
pub fn new() -> Self { Self }
/// True when `name` collides with a hard-coded system command.
pub fn is_reserved(name: &str) -> bool {
RESERVED.contains(&name.to_ascii_lowercase().as_str())
}
/// Expand a command template by substituting the user's arguments.
/// Thin wrapper over [`core_api::command::expand_template`] so callers that hold
/// the concrete manager (e.g. the WS handler) keep their existing call sites.
pub fn expand(&self, template: &str, args: &str) -> String {
expand_template(template, args)
}
}
impl Default for LlmCommandManager {
fn default() -> Self { Self::new() }
}
impl CommandApi for LlmCommandManager {
/// Every enabled, non-reserved command (metadata only — no template body),
/// sorted by name. Tolerant of a missing `commands/` directory (returns empty).
fn list_enabled(&self) -> Vec<CommandInfo> {
let mut out = Vec::new();
let dir = match std::fs::read_dir(COMMANDS_DIR) {
Ok(d) => d,
Err(_) => return out, // no commands/ dir yet → no custom commands
};
for entry in dir.flatten() {
let path = entry.path();
if !path.is_dir() { continue; }
let name = match path.file_name().and_then(|n| n.to_str()) {
Some(n) if !n.is_empty() => n.to_string(),
_ => continue,
};
if Self::is_reserved(&name) { continue; }
if !path.join("meta.json").exists() || !path.join("COMMAND.md").exists() {
continue;
}
let raw = match std::fs::read_to_string(path.join("meta.json"))
.ok()
.and_then(|s| serde_json::from_str::<RawMeta>(&s).ok())
{
Some(r) => r,
None => { warn!(command = %name, "skipping command: missing/invalid meta.json"); continue; }
};
if !raw.enabled { continue; }
out.push(CommandInfo {
name,
description: raw.description,
});
}
out.sort_by(|a, b| a.name.cmp(&b.name));
out
}
/// Resolve a command by name (case-insensitive), loading its template body.
/// Returns `None` when the command does not exist, is disabled, or is reserved.
fn resolve(&self, name: &str) -> Option<ResolvedCommand> {
if Self::is_reserved(name) { return None; }
let name = name.to_ascii_lowercase();
let base = std::path::Path::new(COMMANDS_DIR).join(&name);
let raw: RawMeta = std::fs::read_to_string(base.join("meta.json"))
.ok()
.and_then(|s| serde_json::from_str(&s).ok())?;
if !raw.enabled { return None; }
let template = std::fs::read_to_string(base.join("COMMAND.md")).ok()?;
Some(ResolvedCommand { name, template })
}
}
+512
View File
@@ -0,0 +1,512 @@
//! Context compaction — reduces LLM context size by summarising old messages.
//!
//! # Responsibility
//! [`ContextCompactor`] is a stateless service (all state lives in the DB).
//! It is shared via `Arc` across all [`ChatSessionHandler`]s.
//!
//! It is triggered **at the start of a turn** when the previous turn's
//! `input_tokens` exceeds the configured threshold (Opzione C from the design
//! doc), or manually via `force_compact`. Ephemeral sessions (cron, tic)
//! are always skipped.
//!
//! # Compaction flow
//! ```text
//! handle_message()
//! └─► ContextCompactor::try_compact(pool, stack_id, last_input_tokens)
//! │
//! ├─ guard: tokens < threshold → return Ok(false)
//! ├─ guard: is_ephemeral → return Ok(false)
//! │
//! └─► do_compact(pool, session_id, stack_id, effective_tokens)
//! ├─ load latest summary (if any)
//! ├─ load raw messages since last summary boundary
//! │ (or all messages if no prior summary)
//! ├─ split: to_summarise = messages[0 .. len - keep_recent]
//! │ to_keep_raw = messages[len - keep_recent ..]
//! ├─ if to_summarise is empty → return Ok(false)
//! ├─ build compaction prompt (system hard-coded + user = conversation text)
//! ├─ call LLM (no tools, strength-based AUTO selection)
//! ├─ save summary to chat_summaries
//! └─ publish BusEvent::CompactionDone
//!
//! force_compact() skips the threshold guard and calls do_compact() directly.
//! ```
//!
//! # build_openai_messages after compaction
//! ```text
//! latest_summary = chat_summaries::latest_for_stack(pool, stack_id)
//! if let Some(s) = latest_summary:
//! inject <summary>…</summary> after system prompt
//! load messages with id > s.covers_up_to_message_id
//! else:
//! load all messages (current behaviour)
//! apply max_history_messages drain as safety floor (only when compaction is disabled)
//! ```
use std::sync::Arc;
use serde_json::json;
use sqlx::SqlitePool;
use tracing::{debug, info, warn};
use crate::chat_event_bus::{ChatEventBus, CompactionEvent};
use crate::chatbot::ChatOptions;
use crate::config::CompactionConfig;
use crate::db::{chat_history, chat_llm_tools, chat_summaries};
use crate::llm::LlmManager;
// ── Compaction constants (ported from Hermes context_compressor.py) ──────────
//
// SUMMARY_PREFIX — prepended to every stored summary when injected as context.
// Tells the LLM this is historical reference, not live instructions.
// SUMMARIZER_PREAMBLE — system/user-message preamble for the summarisation LLM call.
// SUMMARY_TEMPLATE — structured section template the LLM must follow.
/// Prefix prepended to the summary content when it is injected into the
/// message array as context for the main agent. Exposed as `pub` so that
/// `build_openai_messages` can use the same wording.
pub const SUMMARY_PREFIX: &str = "\
[CONTEXT COMPACTION — REFERENCE ONLY] Earlier turns were compacted \
into the summary below. This is a handoff from a previous context \
window — treat it as background reference, NOT as active instructions. \
Do NOT answer questions or fulfill requests mentioned in this summary; \
they were already addressed. \
Your current task is identified in the '## Active Task' section of the \
summary — resume exactly from there. \
Your system prompt and any injected memory files are ALWAYS authoritative \
— never deprioritize them due to this compaction note. \
Respond ONLY to the latest user message that appears AFTER this summary. \
The current session state (files, config, etc.) may reflect work \
described here — avoid repeating it:";
/// Preamble shared by both first-compaction and iterative-update prompts.
/// Wording is deliberately plain to avoid content-filter false positives.
const SUMMARIZER_PREAMBLE: &str = "\
You are a summarization agent creating a context checkpoint. \
Treat the conversation turns below as source material for a \
compact record of prior work. \
Produce only the structured summary; do not add a greeting, \
preamble, or prefix. \
Write the summary in the same language the user was using in the \
conversation — do not translate or switch to English. \
NEVER include API keys, tokens, passwords, secrets, credentials, \
or connection strings in the summary — replace any that appear \
with [REDACTED]. Note that the user may have had credentials present, \
but do not preserve their values.";
/// Structured section template the summariser must fill in.
const SUMMARY_TEMPLATE: &str = "\
## Active Task
[THE SINGLE MOST IMPORTANT FIELD. Copy the user's most recent request or \
task assignment verbatim — the exact words they used. If multiple tasks \
were requested and only some are done, list only the ones NOT yet completed. \
Continuation should pick up exactly here. Example: \
\"User asked: 'Now refactor the auth module to use JWT instead of sessions'\" \
If no outstanding task exists, write \"None.\"]
## Goal
[What the user is trying to accomplish overall]
## Constraints & Preferences
[User preferences, coding style, constraints, important decisions]
## Completed Actions
[Numbered list of concrete actions taken — include tool used, target, and outcome.
Format each as: N. ACTION target — outcome [tool: name]
Example:
1. READ config.rs:45 — found == should be != [tool: read_file]
2. EDIT config.rs:45 — changed == to != [tool: write_file]
3. BUILD `cargo build` — succeeded, 0 errors [tool: execute_cmd]
Be specific with file paths, commands, line numbers, and results.]
## Active State
[Current working state — include:
- Working directory and branch (if applicable)
- Modified/created files with brief note on each
- Build/test status
- Any running processes or servers
- Environment details that matter]
## In Progress
[Work currently underway — what was being done when compaction fired]
## Blocked
[Any blockers, errors, or issues not yet resolved. Include exact error messages.]
## Key Decisions
[Important technical decisions and WHY they were made]
## Resolved Questions
[Questions the user asked that were ALREADY answered — include the answer so it is not repeated]
## Pending User Asks
[Questions or requests from the user that have NOT yet been answered or fulfilled. If none, write \"None.\"]
## Relevant Files
[Files read, modified, or created — with brief note on each]
## Remaining Work
[What remains to be done — framed as context, not instructions]
## Critical Context
[Any specific values, error messages, configuration details, or data that would \
be lost without explicit preservation. NEVER include API keys, tokens, passwords, \
or credentials — write [REDACTED] instead.]
Write only the summary body. Do not include any preamble or prefix.";
// ── Public API ────────────────────────────────────────────────────────────────
pub struct ContextCompactor {
config: CompactionConfig,
llm_manager: Arc<LlmManager>,
event_bus: Arc<ChatEventBus>,
}
impl ContextCompactor {
pub fn new(
config: CompactionConfig,
llm_manager: Arc<LlmManager>,
event_bus: Arc<ChatEventBus>,
) -> Self {
Self { config, llm_manager, event_bus }
}
/// Attempt to compact the conversation history for `stack_id`.
///
/// * `last_input_tokens` — input tokens from the **previous** turn.
/// Pass `0` when the provider did not report usage (a character-count
/// estimate is used as fallback in that case).
/// * `is_ephemeral` — skip compaction for short-lived automated sessions.
///
/// Returns `true` if a new summary was written, `false` if skipped.
pub async fn try_compact(
&self,
pool: &SqlitePool,
session_id: i64,
stack_id: i64,
last_input_tokens: u32,
is_ephemeral: bool,
) -> anyhow::Result<bool> {
if is_ephemeral {
return Ok(false);
}
let effective_tokens = if last_input_tokens > 0 {
last_input_tokens
} else {
let est = chat_history::estimate_tokens_for_stack(pool, stack_id).await?;
debug!(stack_id, estimate = est, "compactor: no usage data, using char estimate");
est
};
if effective_tokens < self.config.threshold_tokens {
return Ok(false);
}
info!(
stack_id,
effective_tokens,
threshold = self.config.threshold_tokens,
"compactor: threshold exceeded, starting compaction"
);
self.do_compact(pool, session_id, stack_id, effective_tokens).await
}
/// Force compaction regardless of the token threshold.
/// Still respects the ephemeral guard.
///
/// Returns `true` if a new summary was written, `false` if skipped.
pub async fn force_compact(
&self,
pool: &SqlitePool,
session_id: i64,
stack_id: i64,
is_ephemeral: bool,
) -> anyhow::Result<bool> {
if is_ephemeral {
return Ok(false);
}
let effective_tokens = chat_history::estimate_tokens_for_stack(pool, stack_id).await?;
info!(
stack_id,
effective_tokens,
"compactor: manual compaction triggered"
);
self.do_compact(pool, session_id, stack_id, effective_tokens).await
}
/// Core compaction logic shared by `try_compact` and `force_compact`.
/// Loads messages, splits at the keep_recent boundary, calls the summariser
/// LLM, persists the summary, and publishes a `CompactionDone` event.
async fn do_compact(
&self,
pool: &SqlitePool,
session_id: i64,
stack_id: i64,
effective_tokens: u32,
) -> anyhow::Result<bool> {
let prior_summary = chat_summaries::latest_for_stack(pool, stack_id).await?;
let messages = match &prior_summary {
Some(s) => chat_history::for_stack_since(pool, stack_id, s.covers_up_to_message_id).await?,
None => chat_history::for_stack(pool, stack_id).await?,
};
let keep = self.config.keep_recent;
if messages.len() <= keep {
debug!(
stack_id,
messages = messages.len(),
keep,
"compactor: not enough messages to summarise beyond keep_recent, skipping"
);
return Ok(false);
}
let raw_split = messages.len() - keep;
let split = (0..=raw_split)
.rev()
.find(|&i| {
i == 0 || matches!(
messages[i].role,
chat_history::Role::User | chat_history::Role::Agent
)
})
.unwrap_or(0);
if split == 0 {
debug!(stack_id, "compactor: no suitable split point found, skipping");
return Ok(false);
}
let to_summarise = &messages[..split];
let last_covered_id = to_summarise.last().expect("to_summarise is non-empty").id;
let conversation_text = self
.format_for_summary(pool, to_summarise, prior_summary.as_ref().map(|s| s.content.as_str()))
.await?;
let (client_name, llm) = self.llm_manager
.resolve(None, None, self.config.strength)
.await?;
info!(
stack_id,
client = %client_name,
messages_covered = to_summarise.len(),
last_covered_id,
"compactor: calling LLM for summary"
);
let messages_payload = vec![
json!({ "role": "user", "content": conversation_text }),
];
let options = ChatOptions {
model: llm.model.clone(),
max_tokens: None,
temperature: Some(0.3),
session_id: Some(session_id),
stack_id: Some(stack_id),
};
let turn = llm.client.chat_with_tools(&messages_payload, &[], &options).await
.map_err(|e| {
warn!(stack_id, error = %e, "compactor: LLM call failed");
e
})?;
let summary_text = match turn {
crate::chatbot::LlmTurn::Message(resp) => resp.content,
crate::chatbot::LlmTurn::ToolCalls { content, .. } => {
warn!(stack_id, "compactor: unexpected tool calls in summary response, using content");
content
}
};
if summary_text.trim().is_empty() {
warn!(stack_id, "compactor: LLM returned empty summary, skipping save");
return Ok(false);
}
let summary_id = chat_summaries::save(pool, stack_id, &summary_text, last_covered_id).await?;
info!(
stack_id,
summary_id,
last_covered_id,
"compactor: summary saved"
);
self.event_bus.compaction_done(CompactionEvent {
session_id,
stack_id,
summary_id,
covers_up_to_message_id: last_covered_id,
triggered_by_tokens: effective_tokens,
});
Ok(true)
}
// ── Private helpers ───────────────────────────────────────────────────────
/// Builds the full prompt for the summarisation LLM call (Hermes-style).
///
/// Returns a single string intended to be sent as a `user` message.
/// The preamble, conversation transcript, and structured template are all
/// concatenated, matching how Hermes' `_generate_summary` works.
///
/// * First compaction — `prior_summary` is `None`.
/// * Subsequent compaction — `prior_summary` contains the previous summary body
/// (without `SUMMARY_PREFIX`) so the LLM can produce an updated, non-nested summary.
async fn format_for_summary(
&self,
pool: &SqlitePool,
messages: &[chat_history::ChatMessage],
prior_summary: Option<&str>,
) -> anyhow::Result<String> {
let transcript = self.serialize_for_summary(pool, messages).await?;
let prompt = if let Some(prev) = prior_summary {
format!(
"{SUMMARIZER_PREAMBLE}\n\n\
You are updating a context compaction summary. A previous compaction produced \
the summary below. New conversation turns have occurred since then and need \
to be incorporated.\n\n\
PREVIOUS SUMMARY:\n{prev}\n\n\
NEW TURNS TO INCORPORATE:\n{transcript}\n\n\
Update the summary using this exact structure. PRESERVE all existing information \
that is still relevant. ADD new completed actions to the numbered list (continue \
numbering). Move items from \"In Progress\" to \"Completed Actions\" when done. \
Move answered questions to \"Resolved Questions\". Update \"Active State\" to \
reflect current state. Remove information only if it is clearly obsolete. \
CRITICAL: Update \"## Active Task\" to reflect the user's most recent unfulfilled \
request — this is the most important field for task continuity.\n\n\
{SUMMARY_TEMPLATE}"
)
} else {
format!(
"{SUMMARIZER_PREAMBLE}\n\n\
Create a structured checkpoint summary for the conversation after earlier turns \
are compacted. The summary should preserve enough detail for continuity without \
re-reading the original turns.\n\n\
TURNS TO SUMMARIZE:\n{transcript}\n\n\
Use this exact structure:\n\n\
{SUMMARY_TEMPLATE}"
)
};
Ok(prompt)
}
/// Serialises conversation messages into Hermes-style labeled text for the summariser.
///
/// Format:
/// ```text
/// [USER]: text…
///
/// [ASSISTANT]: text…
/// [Tool calls:
/// tool_name(args…)
/// ]
///
/// [TOOL RESULT tc_N]: result…
/// ```
///
/// Long content is truncated with a head+tail strategy (preserving the start and
/// end of the text) rather than a simple prefix cut.
async fn serialize_for_summary(
&self,
pool: &SqlitePool,
messages: &[chat_history::ChatMessage],
) -> anyhow::Result<String> {
let mut parts: Vec<String> = Vec::new();
for msg in messages {
match msg.role {
chat_history::Role::User | chat_history::Role::Agent => {
let content = truncate_head_tail(msg.content.trim(), 6000, 1500);
parts.push(format!("[USER]: {content}"));
}
chat_history::Role::Assistant => {
let mut content = truncate_head_tail(msg.content.trim(), 6000, 1500);
let tool_calls = chat_llm_tools::for_message(pool, msg.id).await?;
if !tool_calls.is_empty() {
let tc_lines: String = tool_calls
.iter()
.map(|tc| {
let args = tc.arguments.as_deref()
.map(|a| truncate(a, 1200))
.unwrap_or_default();
format!(" {}({})", tc.name, args)
})
.collect::<Vec<_>>()
.join("\n");
content.push_str(&format!("\n[Tool calls:\n{tc_lines}\n]"));
}
parts.push(format!("[ASSISTANT]: {content}"));
// Tool results as separate labeled entries — mirrors Hermes'
// `[TOOL RESULT {call_id}]` entries in the serialised transcript.
for tc in &tool_calls {
let result = match tc.status.as_str() {
"done" => tc.result.as_deref()
.map(|r| truncate_head_tail(r, 4000, 1500))
.unwrap_or_default(),
_ => "(failed or interrupted)".to_string(),
};
parts.push(format!("[TOOL RESULT tc_{}]: {result}", tc.id));
}
}
}
}
Ok(parts.join("\n\n"))
}
}
/// Truncate a string to at most `max_chars`, appending "…" if truncated.
fn truncate(s: &str, max_chars: usize) -> String {
let s = s.trim();
if s.chars().count() <= max_chars {
s.to_string()
} else {
let end = s.char_indices()
.nth(max_chars)
.map(|(i, _)| i)
.unwrap_or(s.len());
format!("{}", &s[..end])
}
}
/// Keep the first `head_chars` and last `tail_chars` of a string, inserting
/// `\n...[truncated]...\n` in the middle when the string is longer than their sum.
///
/// Mirrors Hermes' `_CONTENT_HEAD` + `_CONTENT_TAIL` strategy so the summariser
/// always sees both the beginning context and the ending result of verbose outputs.
fn truncate_head_tail(s: &str, head_chars: usize, tail_chars: usize) -> String {
let s = s.trim();
let char_count = s.chars().count();
let total = head_chars + tail_chars;
if char_count <= total {
return s.to_string();
}
let head_end = s.char_indices()
.nth(head_chars)
.map(|(i, _)| i)
.unwrap_or(s.len());
let tail_start = s.char_indices()
.nth(char_count - tail_chars)
.map(|(i, _)| i)
.unwrap_or(0);
format!("{}\n...[truncated]...\n{}", &s[..head_end], &s[tail_start..])
}
+118
View File
@@ -0,0 +1,118 @@
use serde::Deserialize;
pub use core_api::provider::LlmStrength;
// ── Core config types ─────────────────────────────────────────────────────────
/// LLM runtime settings (clients are managed via LlmManager / DB, not here).
#[derive(Debug, Deserialize)]
pub struct LlmConfig {
pub max_history_messages: usize,
pub max_tool_rounds: Option<usize>,
/// Maximum number of synchronous sub-agents run concurrently when the LLM emits
/// a homogeneous batch of sub-agent calls in one response. Omit to use the
/// default (`DEFAULT_MAX_PARALLEL_SUBAGENTS`). `1` forces sequential dispatch.
#[serde(default)]
pub max_parallel_subagents: Option<usize>,
/// When set, tool results from previous turns that exceed this many characters are
/// replaced at context-build time with a short placeholder. The original result is
/// always preserved in the database (and shown in the frontend); only what the LLM
/// sees in subsequent turns is affected. Omit or set to `null` to disable.
pub max_tool_result_chars: Option<usize>,
/// Request/response logging configuration. Omit or set `enabled: false` to disable.
pub requests_log: Option<LlmRequestsLogConfig>,
/// Context compaction settings. Omit to disable automatic compaction.
pub compaction: Option<CompactionConfig>,
/// Controls how the current date/time is injected into each LLM request.
#[serde(default)]
pub datetime: DatetimeConfig,
}
/// Controls date/time injection in the dynamic tail of each LLM request.
#[derive(Debug, Clone, Deserialize)]
pub struct DatetimeConfig {
/// Inject the current date/time into the LLM context. Default: true.
#[serde(default = "default_true")]
pub enabled: bool,
/// When set, round the injected time down to the nearest N-minute boundary.
pub round_minutes: Option<u32>,
/// IANA timezone name to use when formatting the injected timestamp.
/// Populated at startup from the global `timezone` config field.
#[serde(skip)]
pub timezone: Option<String>,
}
impl Default for DatetimeConfig {
fn default() -> Self {
Self { enabled: true, round_minutes: None, timezone: None }
}
}
/// Context compaction: summarises conversation history when the LLM context
/// exceeds `threshold_tokens`.
#[derive(Debug, Clone, Deserialize)]
pub struct CompactionConfig {
/// Trigger compaction when the previous turn consumed more than this many input tokens.
pub threshold_tokens: u32,
/// Number of recent messages to keep outside the summary. Defaults to 6.
#[serde(default = "default_keep_recent")]
pub keep_recent: usize,
/// Minimum LLM strength to use for generating summaries via AUTO selection.
pub strength: Option<LlmStrength>,
}
/// TIC background event processor settings.
#[derive(Debug, Clone, Deserialize)]
pub struct TicConfig {
/// Interval between ticks, in seconds. Default: 900 (15 minutes).
#[serde(default = "default_tic_interval_secs")]
pub interval_secs: u64,
/// Maximum number of events processed per tick. Default: 50.
#[serde(default = "default_tic_batch_size")]
pub batch_size: i64,
}
impl Default for TicConfig {
fn default() -> Self {
Self { interval_secs: default_tic_interval_secs(), batch_size: default_tic_batch_size() }
}
}
/// Cron scheduler settings.
#[derive(Debug, Default, Deserialize)]
pub struct CronConfig {}
/// Settings for the LLM request/response log (table `llm_requests`).
#[derive(Debug, Clone, Deserialize)]
pub struct LlmRequestsLogConfig {
#[serde(default)]
pub enabled: bool,
#[serde(default = "default_true")]
pub request_payload_save: bool,
#[serde(default = "default_true")]
pub response_payload_save: bool,
#[serde(default = "default_true")]
pub request_header_save: bool,
#[serde(default = "default_true")]
pub response_header_save: bool,
pub cleanup_request_payload_after: Option<u32>,
pub cleanup_response_payload_after: Option<u32>,
pub cleanup_headers_after: Option<u32>,
pub cleanup_rows_after: Option<u32>,
}
fn default_true() -> bool { true }
fn default_keep_recent() -> usize { 6 }
fn default_tic_interval_secs() -> u64 { 900 }
fn default_tic_batch_size() -> i64 { 50 }
// ── CoreConfig ────────────────────────────────────────────────────────────────
/// Core application config — passed to `Skald::new()`.
/// No HTTP/server knowledge. Derived from `Config` via `Config::into_split()`.
pub struct CoreConfig {
pub llm: LlmConfig,
pub tic: TicConfig,
pub cron: CronConfig,
pub timezone: Option<String>,
}
+43
View File
@@ -0,0 +1,43 @@
use std::sync::Arc;
use sqlx::SqlitePool;
pub struct GlobalConfigManager {
pool: Arc<SqlitePool>,
}
impl GlobalConfigManager {
pub fn new(pool: Arc<SqlitePool>) -> Self {
Self { pool }
}
pub async fn get(&self, key: &str) -> anyhow::Result<Option<String>> {
let row = sqlx::query_as::<_, (String,)>("SELECT value FROM config WHERE key = ?")
.bind(key)
.fetch_optional(&*self.pool)
.await?;
Ok(row.map(|(v,)| v))
}
pub async fn set(&self, key: &str, value: &str) -> anyhow::Result<()> {
sqlx::query(
"INSERT INTO config (key, value, updated_at) VALUES (?, ?, datetime('now'))
ON CONFLICT(key) DO UPDATE SET
value = excluded.value,
updated_at = excluded.updated_at",
)
.bind(key)
.bind(value)
.execute(&*self.pool)
.await?;
Ok(())
}
pub async fn remove(&self, key: &str) -> anyhow::Result<()> {
sqlx::query("DELETE FROM config WHERE key = ?")
.bind(key)
.execute(&*self.pool)
.await?;
Ok(())
}
}
+695
View File
@@ -0,0 +1,695 @@
use std::str::FromStr;
use std::sync::Arc;
use anyhow::Result;
use chrono::{DateTime, Local, Utc};
use chrono_tz::Tz;
use cron::Schedule;
use sqlx::SqlitePool;
use tokio::sync::mpsc;
use tokio::time::Duration;
use tracing::{error, info};
use core_api::system_bus::{SystemEvent, SystemEventBus};
use crate::chat_hub::ChatHub;
use crate::db::chat_sessions;
use crate::db::scheduled_jobs::{self, ScheduledJob};
use crate::session::manager::ChatSessionManager;
pub struct TaskManager {
pool: Arc<SqlitePool>,
tz: Option<Tz>,
session: std::sync::OnceLock<Arc<ChatSessionManager>>,
hub: std::sync::OnceLock<Arc<ChatHub>>,
self_arc: std::sync::OnceLock<Arc<Self>>,
system_bus: Arc<SystemEventBus>,
}
/// Returns `(next_utc, is_single)` where `is_single` is `true` when the
/// schedule has no second fire time after the first — i.e. the expression
/// can only ever fire once. Falls back to system local time when `tz` is `None`.
fn next_fire_and_single(schedule: &Schedule, tz: Option<Tz>) -> Option<(DateTime<Utc>, bool)> {
if let Some(tz) = tz {
let mut it = schedule.upcoming(tz);
let first = it.next()?.with_timezone(&Utc);
Some((first, it.next().is_none()))
} else {
let mut it = schedule.upcoming(Local);
let first = it.next()?.with_timezone(&Utc);
Some((first, it.next().is_none()))
}
}
fn next_fire(schedule: &Schedule, tz: Option<Tz>) -> Option<DateTime<Utc>> {
next_fire_and_single(schedule, tz).map(|(dt, _)| dt)
}
impl TaskManager {
pub fn new(pool: Arc<SqlitePool>, tz: Option<Tz>, system_bus: Arc<SystemEventBus>) -> Arc<Self> {
Arc::new(Self {
pool,
tz,
session: std::sync::OnceLock::new(),
hub: std::sync::OnceLock::new(),
self_arc: std::sync::OnceLock::new(),
system_bus,
})
}
/// Called once after ChatSessionManager is built, breaking the circular dep.
pub fn set_session(&self, session: Arc<ChatSessionManager>) {
let _ = self.session.set(session);
}
/// Called once after ChatHub is built. Used for completion notifications.
pub fn set_hub(&self, hub: Arc<ChatHub>) {
let _ = self.hub.set(hub);
}
/// Called once after Arc<Self> is available (in skald.rs after new()).
pub fn set_self_arc(&self, arc: Arc<Self>) {
let _ = self.self_arc.set(arc);
}
fn session(&self) -> Result<&Arc<ChatSessionManager>> {
self.session.get().ok_or_else(|| anyhow::anyhow!("cron: session manager not initialized"))
}
fn self_arc(&self) -> Result<Arc<Self>> {
self.self_arc.get().cloned()
.ok_or_else(|| anyhow::anyhow!("cron: self_arc not initialized"))
}
/// Start the background loops. Must be called after set_session().
/// Returns join handles so the caller can await them during shutdown.
pub fn start(self: Arc<Self>, shutdown: tokio_util::sync::CancellationToken) -> Vec<tokio::task::JoinHandle<()>> {
// Main scheduler loop.
let me = Arc::clone(&self);
let sd1 = shutdown.clone();
let h1 = tokio::spawn(async move {
if let Err(e) = me.recover_interrupted().await {
error!("cron: startup recovery failed: {e}");
}
let mut interval = tokio::time::interval(Duration::from_secs(30));
loop {
tokio::select! {
_ = sd1.cancelled() => { info!("cron: scheduler loop stopping"); break; }
_ = interval.tick() => {
if let Err(e) = me.tick().await {
error!("cron tick error: {e}");
}
}
}
}
});
// Cleanup loop: removes single_run jobs completed more than 7 days ago.
let pool = Arc::clone(&self.pool);
let sd2 = shutdown.clone();
let h2 = tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(15)).await;
let mut interval = tokio::time::interval(Duration::from_secs(3600));
loop {
tokio::select! {
_ = sd2.cancelled() => { info!("cron: cleanup loop stopping"); break; }
_ = interval.tick() => {
if let Err(e) = cleanup_expired_single_runs(&pool).await {
error!("cron: cleanup error: {e}");
}
}
}
}
});
vec![h1, h2]
}
async fn recover_interrupted(&self) -> Result<()> {
let session = self.session()?;
let self_arc = self.self_arc()?;
let jobs = scheduled_jobs::list_interrupted(&self.pool).await?;
if jobs.is_empty() { return Ok(()); }
info!("cron: recovering {} interrupted job(s)", jobs.len());
for job in jobs {
let pool = Arc::clone(&self.pool);
let session = Arc::clone(session);
let hub = self.hub.get().cloned();
let task_mgr = Arc::clone(&self_arc);
let tz = self.tz;
tokio::spawn(async move {
if let Err(e) = run_job(&pool, &session, &task_mgr, hub.as_ref(), &job, tz).await {
error!("cron: recovery of job {} ('{}') failed: {e}", job.id, job.title);
}
});
}
Ok(())
}
async fn tick(&self) -> Result<()> {
let session = self.session()?;
let self_arc = self.self_arc()?;
let now = Utc::now().to_rfc3339();
let jobs = scheduled_jobs::list_due(&self.pool, &now).await?;
for job in jobs {
let pool = Arc::clone(&self.pool);
let session = Arc::clone(session);
let hub = self.hub.get().cloned();
let task_mgr = Arc::clone(&self_arc);
let job = job.clone();
let tz = self.tz;
tokio::spawn(async move {
if let Err(e) = run_job(&pool, &session, &task_mgr, hub.as_ref(), &job, tz).await {
error!("cron job {} ('{}') failed: {e}", job.id, job.title);
}
});
}
Ok(())
}
// ── Sync wrappers (called from LLM tools via block_in_place) ─────────────
pub fn list_jobs(&self) -> Result<Vec<ScheduledJob>> {
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current()
.block_on(scheduled_jobs::list(&self.pool))
})
}
/// Validate that `agent_id` names a runnable task agent (non-empty, exists,
/// `type == Task`). Single gate shared by every job-creation entry point so
/// cron / sync / async / project-ticket paths all agree — no silent default.
fn require_task_agent(agent_id: &str) -> Result<()> {
if agent_id.trim().is_empty() {
anyhow::bail!("agent_id is required — specify which task agent runs this task (no default)");
}
crate::agents::load_task_meta(agent_id)?;
Ok(())
}
pub fn add_job(
&self,
title: &str,
description: &str,
cron: &str,
prompt: &str,
agent_id: &str,
single_run: bool,
kind: &str,
parent_session_id: Option<i64>,
run_context: Option<&str>,
) -> Result<ScheduledJob> {
Self::require_task_agent(agent_id)?;
let (first_fire, _is_single, single_run) = if kind == "sync" || kind == "immediate" {
(None, true, true)
} else {
let schedule = Schedule::from_str(cron).map_err(|_| {
anyhow::anyhow!(
"Invalid cron expression: '{cron}'. Use 7-field format: \
sec min hour dom month dow year (e.g. '0 0 9 * * * *' = every day at 9:00)"
)
})?;
let (first, single) = next_fire_and_single(&schedule, self.tz)
.ok_or_else(|| anyhow::anyhow!("Cron expression '{cron}' has no upcoming fire times"))?;
let single_run = single_run || single;
(Some(first.to_rfc3339()), single, single_run)
};
let next_run_at: Option<&str> = first_fire.as_deref();
let job = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(scheduled_jobs::create(
&self.pool, title, description, cron, prompt, agent_id,
single_run, next_run_at, kind, parent_session_id, run_context, None,
))
})?;
Ok(job)
}
/// Execute a task synchronously: creates the DB record, runs it inline,
/// and returns the agent's final response. Blocks until completion.
pub fn add_job_sync(
&self,
title: &str,
description: &str,
prompt: &str,
agent_id: &str,
run_context: Option<&str>,
) -> Result<String> {
Self::require_task_agent(agent_id)?;
let job = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(scheduled_jobs::create(
&self.pool, title, description, "", prompt, agent_id,
true, None, "sync", None, run_context, None,
))
})?;
let session = self.session()?;
let self_arc = self.self_arc()?;
let result = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(
run_job(&self.pool, session, &self_arc, self.hub.get(), &job, self.tz)
)
})?;
Ok(result.unwrap_or_else(|| "(no output)".to_string()))
}
/// Start a task asynchronously: creates the DB record, spawns the run,
/// returns immediately. Result is injected into parent_session_id when done.
pub fn add_job_async(
&self,
title: &str,
description: &str,
prompt: &str,
agent_id: &str,
parent_session_id: i64,
run_context: Option<&str>,
) -> Result<ScheduledJob> {
Self::require_task_agent(agent_id)?;
let job = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(scheduled_jobs::create(
&self.pool, title, description, "", prompt, agent_id,
true, None, "async", Some(parent_session_id), run_context, None,
))
})?;
let pool = Arc::clone(&self.pool);
let session = self.session()?.clone();
let hub = self.hub.get().cloned();
let task_mgr = self.self_arc()?;
let tz = self.tz;
let job_c = job.clone();
tokio::spawn(async move {
if let Err(e) = run_job(&pool, &session, &task_mgr, hub.as_ref(), &job_c, tz).await {
error!("async task {} ('{}') failed: {e}", job_c.id, job_c.title);
}
});
Ok(job)
}
/// Create and immediately spawn an async job with an opaque `origin_ref`.
/// Returns the created `ScheduledJob` (caller uses its `id` for tracking).
/// Unlike `add_job_async`, no `parent_session_id` is set — completion is
/// delivered via `SystemEvent::JobCompleted` on the system bus.
pub fn spawn_async_job(
&self,
title: &str,
description: &str,
prompt: &str,
agent_id: &str,
run_context: Option<&str>,
origin_ref: &str,
) -> Result<scheduled_jobs::ScheduledJob> {
Self::require_task_agent(agent_id)?;
let job = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(scheduled_jobs::create(
&self.pool, title, description, "", prompt, agent_id,
true, None, "async", None, run_context, Some(origin_ref),
))
})?;
let pool = Arc::clone(&self.pool);
let session = self.session()?.clone();
let hub = self.hub.get().cloned();
let task_mgr = self.self_arc()?;
let tz = self.tz;
let job_c = job.clone();
tokio::spawn(async move {
if let Err(e) = run_job(&pool, &session, &task_mgr, hub.as_ref(), &job_c, tz).await {
error!("project-ticket job {} failed: {e}", job_c.id);
}
});
Ok(job)
}
pub fn delete_job(&self, id: i64) -> Result<bool> {
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current()
.block_on(scheduled_jobs::delete(&self.pool, id))
})
}
pub fn toggle_job(&self, id: i64, enabled: bool) -> Result<bool> {
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async {
let found = scheduled_jobs::set_enabled(&self.pool, id, enabled).await?;
if found && enabled {
// Recalculate next_run_at when re-enabling so a stale timestamp
// doesn't cause an immediate spurious fire.
let jobs = scheduled_jobs::list(&self.pool).await?;
if let Some(job) = jobs.iter().find(|j| j.id == id) {
let tz = self.tz;
if let Some(next) = Schedule::from_str(&job.cron)
.ok()
.and_then(|s| next_fire(&s, tz))
.map(|t| t.to_rfc3339())
{
scheduled_jobs::set_next_run_at(&self.pool, id, &next).await?;
}
}
}
Ok(found)
})
})
}
}
// ── Job execution ─────────────────────────────────────────────────────────────
async fn run_job(
pool: &SqlitePool,
session: &ChatSessionManager,
task_mgr: &Arc<TaskManager>,
hub: Option<&Arc<ChatHub>>,
job: &ScheduledJob,
tz: Option<Tz>,
) -> Result<Option<String>> {
info!("running {} task {} ('{}')", job.kind, job.id, job.title);
let started_at = Utc::now();
let (session_id, _) = session.create_session(&job.agent_id, "cron", false, true, None).await?;
scheduled_jobs::set_running(pool, job.id, session_id).await?;
if let Some(rc) = &job.run_context {
chat_sessions::set_run_context(pool, session_id, Some(rc.as_str())).await.ok();
}
let handler = session.get_or_create_handler(session_id).await?;
handler.set_context_label(format!("CronJob: {}", job.title));
if job.kind == "async" {
if let Some(parent_id) = job.parent_session_id {
handler.set_scratchpad_session_id(parent_id);
}
}
let job_context = format!(
"[Job context]\nJob ID: {}{}\nTime: {} UTC",
job.id, job.title,
started_at.format("%Y-%m-%d %H:%M"),
);
// Build interface_tools: execute_subtask for background sessions (sync only, no async/cron).
let task_mgr_clone = Arc::clone(task_mgr);
let execute_subtask_tool = build_execute_subtask_tool(task_mgr_clone, job.run_context.clone());
// Use a large buffer and drain rx concurrently with handle_message to avoid
// deadlock: handle_message may emit many events (ToolStart/ToolDone/Thinking
// per tool call), and the channel blocks when full if nobody is reading.
let (tx, mut rx) = mpsc::channel(512);
let handler_arc = Arc::clone(&handler);
let prompt = job.prompt.clone();
let ctx = job_context.clone();
let jh = tokio::spawn(async move {
handler_arc.handle_message(
&prompt,
None,
None,
Some(ctx),
None,
vec![execute_subtask_tool],
std::collections::HashMap::new(),
tx,
false,
None,
None, // non-interactive: no live user-message injection
).await
});
// Drain events concurrently. rx closes when the last tx clone is dropped,
// which happens only after resume_turn() completes the full sub-agent chain.
while let Some(_) = rx.recv().await {}
let handle_result = jh.await
.unwrap_or_else(|e| Err(anyhow::anyhow!("run_job task panicked: {e}")));
let completed_at = Utc::now();
let duration_ms = (completed_at - started_at).num_milliseconds();
let final_response = last_assistant_message(pool, session_id).await.ok().flatten();
let next_run_at: Option<String> = if job.single_run || job.kind != "cron" {
None
} else {
Schedule::from_str(&job.cron).ok()
.and_then(|s| next_fire(&s, tz))
.map(|t| t.to_rfc3339())
};
match handle_result {
Ok(_) => {
record_job_run(pool, job.id, session_id, &started_at.to_rfc3339(),
&completed_at.to_rfc3339(), duration_ms,
"completed", final_response.as_deref(), None).await?;
scheduled_jobs::finish_run(pool, job.id, next_run_at.as_deref()).await?;
task_mgr.system_bus.send(SystemEvent::JobCompleted {
job_id: job.id,
origin_ref: job.origin_ref.clone(),
result: final_response.clone(),
error: None,
});
match job.kind.as_str() {
"cron" => {
if let Some(hub) = hub {
let outcome = final_response.as_deref().unwrap_or("(no output)");
hub.notify(crate::notification::Notification {
source: "cron".into(),
event_type: "cron_result".into(),
summary: format!(
"Cron job \"{}\" (ID {}) completed: {}",
job.title, job.id, outcome,
),
event_time: Utc::now().to_rfc3339(),
refs: serde_json::json!({ "job_id": job.id, "title": job.title }),
}).await.ok();
}
}
"async" => {
if let Some(parent_id) = job.parent_session_id {
if let Some(hub) = hub {
inject_async_result(
pool,
hub,
parent_id,
job.id,
&job.title,
final_response.as_deref().unwrap_or("(no output)"),
).await;
}
}
}
_ => {} // sync: result was already returned inline via add_job_sync
}
info!("{} task {} done", job.kind, job.id);
Ok(final_response)
}
Err(e) => {
let err_str = e.to_string();
record_job_run(pool, job.id, session_id, &started_at.to_rfc3339(),
&completed_at.to_rfc3339(), duration_ms,
"failed", None, Some(&err_str)).await?;
scheduled_jobs::finish_run(pool, job.id, next_run_at.as_deref()).await?;
task_mgr.system_bus.send(SystemEvent::JobCompleted {
job_id: job.id,
origin_ref: job.origin_ref.clone(),
result: None,
error: Some(err_str.clone()),
});
if let Some(hub) = hub {
hub.notify(crate::notification::Notification {
source: "cron".into(),
event_type: "cron_error".into(),
summary: format!(
"Cron job \"{}\" (ID {}) failed: {} (check the logs)",
job.title, job.id, err_str,
),
event_time: Utc::now().to_rfc3339(),
refs: serde_json::json!({ "job_id": job.id, "title": job.title }),
}).await.ok();
}
Err(e)
}
}
}
/// Injects an async task result into the parent session using the same pattern as
/// the notification system: writes a synthetic assistant message + completed
/// `task_completed` tool call directly to the DB, then calls `hub.resume()` so
/// the parent LLM wakes up and events are properly bridged to the WebSocket.
async fn inject_async_result(
pool: &SqlitePool,
hub: &Arc<ChatHub>,
parent_session_id: i64,
task_id: i64,
task_title: &str,
result: &str,
) {
// Resolve source_id from the parent session row.
let source_id = match crate::db::chat_sessions::find_by_id(pool, parent_session_id).await {
Ok(Some(s)) => s.source,
Ok(None) => { error!("inject_async_result: session {parent_session_id} not found"); return; }
Err(e) => { error!("inject_async_result: DB error: {e}"); return; }
};
// Get the active stack for the parent session.
let stack = match crate::db::chat_sessions_stack::active_for_session(pool, parent_session_id).await {
Ok(Some(s)) => s,
Ok(None) => { error!("inject_async_result: no active stack for session {parent_session_id}"); return; }
Err(e) => { error!("inject_async_result: stack lookup failed: {e}"); return; }
};
// Write a synthetic assistant message (reasoning trace).
let reasoning = format!(
"The system is notifying me that async task #{task_id} ('{}') has completed. \
Let me process the result via task_completed.",
task_title,
);
let assistant_id = match crate::db::chat_history::append(
pool, stack.id, &crate::db::chat_history::Role::Assistant,
"", true, Some(&reasoning),
).await {
Ok(id) => id,
Err(e) => { error!("inject_async_result: append assistant failed: {e}"); return; }
};
// Write the completed task_completed tool call with the result payload.
let result_json = serde_json::to_string(&serde_json::json!({
"task_id": task_id,
"title": task_title,
"result": result,
})).unwrap_or_else(|_| "{}".to_string());
let tool_call_id = match crate::db::chat_llm_tools::append(
pool, assistant_id, "task_completed",
&serde_json::json!({"task_id": task_id}).to_string(),
).await {
Ok(id) => id,
Err(e) => { error!("inject_async_result: append tool call failed: {e}"); return; }
};
if let Err(e) = crate::db::chat_llm_tools::complete(pool, tool_call_id, &result_json, "string").await {
error!("inject_async_result: complete tool call failed: {e}"); return;
}
info!(parent_session_id, task_id, task_title, "inject_async_result: resuming parent session");
if let Err(e) = hub.resume(&source_id).await {
error!("inject_async_result: hub.resume failed: {e}");
}
}
/// Builds the `execute_subtask` InterfaceTool injected into background sessions.
/// Background tasks can only run synchronous sub-tasks — no cron or async.
fn build_execute_subtask_tool(task_mgr: Arc<TaskManager>, run_context: Option<String>) -> crate::session::handler::InterfaceTool {
use crate::session::handler::{InterfaceTool, ToolFuture};
use serde_json::json;
InterfaceTool {
definition: json!({
"type": "function",
"function": {
"name": crate::tools::tool_names::EXECUTE_SUBTASK,
"description": "Run a synchronous sub-task and return its result. Blocks until the sub-task completes.",
"parameters": {
"type": "object",
"required": ["title", "prompt", "agent_id"],
"properties": {
"title": { "type": "string", "description": "Short name for this sub-task" },
"description": { "type": "string", "description": "What this sub-task does" },
"prompt": { "type": "string", "description": "Prompt sent to the agent" },
"agent_id": { "type": "string", "description": "Task agent to run (required; e.g. software-engineer, researcher, generalist)" }
}
}
}
}),
handler: Arc::new(move |args: serde_json::Value| -> ToolFuture {
let tm = Arc::clone(&task_mgr);
let title = args["title"].as_str().unwrap_or("").to_string();
let desc = args["description"].as_str().unwrap_or("").to_string();
let prompt = args["prompt"].as_str().unwrap_or("").to_string();
let agent_id = args["agent_id"].as_str().unwrap_or("").to_string();
let run_context = run_context.clone();
Box::pin(async move {
tokio::task::spawn_blocking(move || {
tm.add_job_sync(&title, &desc, &prompt, &agent_id, run_context.as_deref())
})
.await
.map_err(|e| anyhow::anyhow!("execute_subtask task panicked: {e}"))?
})
}),
}
}
// ── Helpers ───────────────────────────────────────────────────────────────────
async fn record_job_run(
pool: &SqlitePool,
job_id: i64,
session_id: i64,
started_at: &str,
completed_at: &str,
duration_ms: i64,
status: &str,
response: Option<&str>,
error: Option<&str>,
) -> Result<()> {
crate::db::job_runs::insert(
pool, job_id, Some(session_id),
started_at, completed_at, duration_ms,
status, response, error,
).await.map(|_| ())
}
/// Returns the most recent successful assistant message in the given session.
async fn last_assistant_message(pool: &SqlitePool, session_id: i64) -> Result<Option<String>> {
let row: Option<(String,)> = sqlx::query_as(
"SELECT ch.content
FROM chat_history ch
JOIN chat_sessions_stack css ON ch.session_stack_id = css.id
WHERE css.session_id = ? AND ch.role = 'assistant' AND ch.status = 'ok'
ORDER BY ch.id DESC
LIMIT 1",
)
.bind(session_id)
.fetch_optional(pool)
.await?;
Ok(row.map(|(c,)| c))
}
async fn cleanup_expired_single_runs(pool: &SqlitePool) -> Result<()> {
// The set of jobs about to be deleted, reused by each cascade step below.
const EXPIRED: &str = "SELECT id FROM scheduled_jobs
WHERE single_run = 1
AND enabled = 0
AND last_run_at < datetime('now', '-7 days')";
// Clear the soft back-reference from project_tickets first: its job_id FK has
// no ON DELETE action, so a ticket still pointing at an expired runner job
// would block the DELETE below with a FOREIGN KEY constraint failure. The
// ticket keeps its result/error — only the (now-GC'd) job pointer is dropped.
sqlx::query(sqlx::AssertSqlSafe(format!(
"UPDATE project_tickets SET job_id = NULL WHERE job_id IN ({EXPIRED})"
)))
.execute(pool)
.await?;
sqlx::query(sqlx::AssertSqlSafe(format!(
"DELETE FROM job_runs WHERE job_id IN ({EXPIRED})"
)))
.execute(pool)
.await?;
let n = sqlx::query(
"DELETE FROM scheduled_jobs
WHERE single_run = 1
AND enabled = 0
AND last_run_at < datetime('now', '-7 days')",
)
.execute(pool)
.await?
.rows_affected();
if n > 0 {
info!("cron: removed {n} expired single-run job(s)");
}
Ok(())
}
+363
View File
@@ -0,0 +1,363 @@
//! Envelope encryption for the per-user databases (blueprint §4 / §5.1).
//!
//! A user's database is encrypted by SQLCipher under a random 256-bit **DEK**.
//! The DEK is never stored in the clear: `users.database_password` holds it
//! sealed with AES-256-GCM under a **KEK** derived from the password with
//! Argon2id.
//!
//! That seal *is* the password verifier. Opening it either yields the DEK — in
//! which case the password was right — or fails the AEAD tag, which means a
//! wrong password, cleanly distinct from a corrupt file. Storing a second hash
//! of the same password beside the seal would only hand an offline attacker an
//! easier target than the seal itself, so encrypted users have no
//! `password_hash` at all.
//!
//! Cleartext users have no database key to bind a verifier to, so they store the
//! Argon2id output directly and it is compared in constant time. Its
//! crackability is harmless: their database is readable by the box owner by
//! design.
//!
//! Nothing here invents a construction — it composes Argon2id, AES-256-GCM and a
//! CSPRNG. Changing the password re-seals the same DEK; the database itself is
//! never re-encrypted.
use std::fmt::{self, Write as _};
use std::sync::LazyLock;
use aes_gcm::aead::Aead;
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
use anyhow::{Context, Result, anyhow, bail};
use argon2::{Algorithm, Argon2, Params, Version};
use rand::Rng as _;
use serde::{Deserialize, Serialize};
use subtle::ConstantTimeEq;
use tokio::sync::Semaphore;
use zeroize::{Zeroize, ZeroizeOnDrop};
/// 256-bit keys throughout: AES-256-GCM for the seal, raw SQLCipher key for the
/// database.
const KEY_LEN: usize = 32;
const NONCE_LEN: usize = 12;
const TAG_LEN: usize = 16;
/// Sealed layout: `nonce ‖ ciphertext ‖ tag`.
const SEALED_LEN: usize = NONCE_LEN + KEY_LEN + TAG_LEN;
pub const SALT_LEN: usize = 16;
/// The only algorithm we accept. Stored per row so it can change later without a
/// migration, but a row asking for anything else is a bug, not a fallback.
const ALGO: &str = "argon2id";
/// Argon2id at 256 MiB is a memory bomb if it runs unbounded: every concurrent
/// login would allocate its own arena. Two at a time keeps a login responsive
/// while capping the peak at ~2× `KdfParams::m`.
///
/// Deliberately a module-level invariant rather than a caller's responsibility —
/// a forgotten permit is a memory incident, not a style problem.
static KDF_GATE: LazyLock<Semaphore> = LazyLock::new(|| Semaphore::new(2));
// ── KDF parameters ────────────────────────────────────────────────────────────
/// Serialized into `users.kdf_params`. Not secret: calibrated on the box when the
/// user is created, and kept per row so raising the cost later needs no migration.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct KdfParams {
pub algo: String,
/// Memory cost, in KiB.
pub m: u32,
/// Time cost (passes).
pub t: u32,
/// Parallelism (lanes).
pub p: u32,
}
impl Default for KdfParams {
/// 256 MiB, ~1s on the weakest box we target.
///
/// Between the OWASP baseline (64 MiB) and the 512 MiB1 GiB of §5.1: four
/// times the cost for an attacker with a GPU, while two concurrent logins
/// peak at 512 MiB rather than 2 GiB and never reach swap.
fn default() -> Self {
Self { algo: ALGO.to_string(), m: 262_144, t: 3, p: 1 }
}
}
impl KdfParams {
pub fn to_json(&self) -> Result<String> {
serde_json::to_string(self).context("serializing kdf_params")
}
pub fn from_json(s: &str) -> Result<Self> {
let p: Self = serde_json::from_str(s).context("parsing kdf_params")?;
if p.algo != ALGO {
bail!("unsupported kdf algorithm: {}", p.algo);
}
Ok(p)
}
/// Cheap parameters, tests only: the real ones cost ~1s and 256 MiB per
/// derivation, which a test suite pays on every register and every login.
#[cfg(test)]
pub fn fast() -> Self {
Self { algo: ALGO.to_string(), m: 64, t: 1, p: 1 }
}
}
// ── Keys ──────────────────────────────────────────────────────────────────────
/// Data Encryption Key: the raw SQLCipher key for one user's database. Random,
/// never derived from the password, so changing the password does not re-encrypt
/// anything.
#[derive(Clone, Zeroize, ZeroizeOnDrop)]
pub struct Dek([u8; KEY_LEN]);
/// Key Encryption Key: `Argon2id(password, salt)`. Seals the [`Dek`], and for a
/// cleartext user its raw bytes *are* the stored verifier.
#[derive(Clone, Zeroize, ZeroizeOnDrop)]
pub struct Kek([u8; KEY_LEN]);
impl Dek {
pub fn random() -> Self {
let mut k = [0u8; KEY_LEN];
rand::rng().fill_bytes(&mut k);
Self(k)
}
/// The `PRAGMA key` value for a raw 256-bit key.
///
/// The double quotes are **part of the value**: sqlx pastes it verbatim into
/// `PRAGMA key = {value};`, and SQLCipher only treats the argument as a raw
/// key when it parses as the blob literal `x'…'`. Given anything else it
/// runs its own KDF over the bytes instead, silently deriving a *different*
/// key from our hex digits — and the database would open, just not the one
/// we meant.
///
/// The returned string is key material. It must never be logged, and it
/// outlives this call inside the pool's connect options.
pub fn to_pragma(&self) -> String {
// `"x'` + 64 hex digits + `'"`
let mut s = String::with_capacity(5 + 2 * KEY_LEN);
s.push_str("\"x'");
for b in self.0 {
let _ = write!(s, "{b:02x}");
}
s.push_str("'\"");
s
}
}
impl Kek {
/// The verifier stored in `users.password_hash` for a cleartext user.
pub fn as_verifier(&self) -> &[u8] {
&self.0
}
}
// Hand-written so a stray `{:?}` — a tracing span, an error context, a panic —
// cannot print key material. Same reasoning as `db::users::Credentials`.
impl fmt::Debug for Dek {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("Dek(<redacted>)")
}
}
impl fmt::Debug for Kek {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("Kek(<redacted>)")
}
}
pub fn random_salt() -> Vec<u8> {
let mut s = vec![0u8; SALT_LEN];
rand::rng().fill_bytes(&mut s);
s
}
// ── Derivation ────────────────────────────────────────────────────────────────
/// `Argon2id(password, salt)` → 256-bit key.
///
/// Runs on a blocking thread — it burns a core and 256 MiB for about a second,
/// which would stall a tokio worker — and behind [`KDF_GATE`].
pub async fn derive_kek(password: &str, salt: &[u8], params: &KdfParams) -> Result<Kek> {
let _permit = KDF_GATE.acquire().await.context("kdf gate closed")?;
let mut password = password.to_owned();
let salt = salt.to_vec();
let params = params.clone();
let out = tokio::task::spawn_blocking(move || {
let result = derive_blocking(password.as_bytes(), &salt, &params);
password.zeroize();
result
})
.await
.context("argon2 task panicked")?;
out
}
fn derive_blocking(password: &[u8], salt: &[u8], params: &KdfParams) -> Result<Kek> {
if params.algo != ALGO {
bail!("unsupported kdf algorithm: {}", params.algo);
}
let p = Params::new(params.m, params.t, params.p, Some(KEY_LEN))
.map_err(|e| anyhow!("invalid argon2 params: {e}"))?;
let mut key = [0u8; KEY_LEN];
Argon2::new(Algorithm::Argon2id, Version::V0x13, p)
.hash_password_into(password, salt, &mut key)
.map_err(|e| anyhow!("argon2 derivation failed: {e}"))?;
Ok(Kek(key))
}
/// Constant-time comparison of a freshly derived key against a stored verifier.
/// `subtle` returns "not equal" for a length mismatch rather than short-circuiting.
pub fn verify(derived: &Kek, stored: &[u8]) -> bool {
derived.0.ct_eq(stored).into()
}
// ── Envelope ──────────────────────────────────────────────────────────────────
/// Why a seal did not open. The distinction is the whole point: a failed tag is
/// an authentication answer, a malformed blob is a broken row.
#[derive(Debug, PartialEq, Eq)]
pub enum KeyError {
/// The AEAD tag did not verify. This *is* the wrong-password signal.
WrongPassword,
Malformed(&'static str),
}
impl fmt::Display for KeyError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
KeyError::WrongPassword => f.write_str("wrong password"),
KeyError::Malformed(w) => write!(f, "malformed wrapped key: {w}"),
}
}
}
impl std::error::Error for KeyError {}
/// Seals `dek` under `kek`. Output is `nonce ‖ ciphertext ‖ tag`, stored as-is in
/// `users.database_password`.
pub fn wrap_dek(kek: &Kek, dek: &Dek) -> Result<Vec<u8>> {
let cipher = Aes256Gcm::new_from_slice(&kek.0).map_err(|e| anyhow!("bad kek length: {e}"))?;
let mut nonce = [0u8; NONCE_LEN];
rand::rng().fill_bytes(&mut nonce);
let sealed = cipher
.encrypt(Nonce::from_slice(&nonce), dek.0.as_slice())
.map_err(|_| anyhow!("AEAD seal failed"))?;
let mut out = Vec::with_capacity(SEALED_LEN);
out.extend_from_slice(&nonce);
out.extend_from_slice(&sealed);
Ok(out)
}
/// Opens a seal produced by [`wrap_dek`]. A failed tag returns
/// [`KeyError::WrongPassword`] — one Argon2id pass answers "is the password
/// right" *and* hands back the key.
pub fn unwrap_dek(kek: &Kek, sealed: &[u8]) -> Result<Dek, KeyError> {
if sealed.len() != SEALED_LEN {
return Err(KeyError::Malformed("unexpected length"));
}
let (nonce, body) = sealed.split_at(NONCE_LEN);
let cipher =
Aes256Gcm::new_from_slice(&kek.0).map_err(|_| KeyError::Malformed("bad kek length"))?;
let mut plain = cipher
.decrypt(Nonce::from_slice(nonce), body)
.map_err(|_| KeyError::WrongPassword)?;
if plain.len() != KEY_LEN {
plain.zeroize();
return Err(KeyError::Malformed("unexpected plaintext length"));
}
let mut dek = [0u8; KEY_LEN];
dek.copy_from_slice(&plain);
plain.zeroize();
Ok(Dek(dek))
}
#[cfg(test)]
mod tests {
use super::*;
async fn kek(password: &str, salt: &[u8]) -> Kek {
derive_kek(password, salt, &KdfParams::fast()).await.unwrap()
}
#[tokio::test]
async fn seal_round_trips_and_the_tag_is_the_password_check() {
let salt = random_salt();
let dek = Dek::random();
let sealed = wrap_dek(&kek("correct horse", &salt).await, &dek).unwrap();
assert_eq!(sealed.len(), SEALED_LEN);
let opened = unwrap_dek(&kek("correct horse", &salt).await, &sealed).unwrap();
assert_eq!(opened.0, dek.0, "the same DEK must come back out");
// The whole auth story: a wrong password is a failed AEAD tag, and it is
// reported as such — not as a corrupt blob.
let err = unwrap_dek(&kek("wrong horse", &salt).await, &sealed).unwrap_err();
assert_eq!(err, KeyError::WrongPassword);
}
#[tokio::test]
async fn a_truncated_seal_is_malformed_not_a_wrong_password() {
let salt = random_salt();
let sealed = wrap_dek(&kek("pw", &salt).await, &Dek::random()).unwrap();
let err = unwrap_dek(&kek("pw", &salt).await, &sealed[..SEALED_LEN - 1]).unwrap_err();
assert!(matches!(err, KeyError::Malformed(_)), "got {err:?}");
}
#[tokio::test]
async fn the_salt_separates_identical_passwords() {
let a = kek("same", &random_salt()).await;
let b = kek("same", &random_salt()).await;
assert_ne!(a.0, b.0, "a per-user salt must defuse rainbow tables");
}
#[tokio::test]
async fn verify_accepts_the_right_password_and_rejects_a_short_verifier() {
let salt = random_salt();
let stored = kek("pw", &salt).await;
assert!(verify(&kek("pw", &salt).await, stored.as_verifier()));
assert!(!verify(&kek("nope", &salt).await, stored.as_verifier()));
// A length mismatch must not panic, and must not compare equal.
assert!(!verify(&kek("pw", &salt).await, &stored.as_verifier()[..16]));
}
#[test]
fn pragma_is_a_quoted_blob_literal() {
let dek = Dek([0xAB; KEY_LEN]);
let p = dek.to_pragma();
assert!(p.starts_with("\"x'") && p.ends_with("'\""), "got {p}");
assert_eq!(p.len(), 5 + 2 * KEY_LEN, "\"x' + 64 hex digits + '\"");
assert!(p.contains("abab"), "lowercase hex");
}
#[test]
fn debug_never_prints_key_material() {
let dek = Dek([0xDE; KEY_LEN]);
let kek = Kek([0xAD; KEY_LEN]);
assert_eq!(format!("{dek:?}"), "Dek(<redacted>)");
assert_eq!(format!("{kek:?}"), "Kek(<redacted>)");
assert!(!format!("{dek:?}").contains("dede"), "no raw key bytes");
assert!(!format!("{kek:?}").contains("adad"), "no raw key bytes");
}
#[test]
fn kdf_params_round_trip_and_reject_foreign_algorithms() {
let p = KdfParams::default();
assert_eq!(p.m, 262_144);
assert_eq!(KdfParams::from_json(&p.to_json().unwrap()).unwrap(), p);
let pbkdf2 = r#"{"algo":"pbkdf2","m":1,"t":1,"p":1}"#;
assert!(KdfParams::from_json(pbkdf2).is_err(), "only argon2id is accepted");
}
}
@@ -0,0 +1,97 @@
use anyhow::Result;
use sqlx::SqlitePool;
use crate::approval::{ApprovalRule, NewApprovalRule, RuleAction};
type RawRow = (i64, Option<String>, Option<String>, String, Option<String>, String, Option<String>, i64, Option<String>);
fn from_raw((id, agent_id, source, tool_pattern, path_pattern, action, note, priority, group_id): RawRow)
-> anyhow::Result<ApprovalRule>
{
let action: RuleAction = action.parse()?;
Ok(ApprovalRule { id, agent_id, source, tool_pattern, path_pattern, action, note, priority, group_id })
}
/// Returns all rules ordered by priority ASC (lowest number = evaluated first).
pub async fn list(pool: &SqlitePool) -> Result<Vec<ApprovalRule>> {
let rows = sqlx::query_as::<_, RawRow>(
"SELECT id, agent_id, source, tool_pattern, path_pattern, action, note, priority, group_id
FROM approval_rules
ORDER BY priority ASC, id ASC",
)
.fetch_all(pool)
.await?;
rows.into_iter().map(from_raw).collect()
}
/// Returns rules applicable to `group_id`: group-specific first, then 'default' as fallback.
/// If `group_id` is `None` or equals `"default"`, only default rules are returned.
pub async fn list_for_group(pool: &SqlitePool, group_id: Option<&str>) -> Result<Vec<ApprovalRule>> {
let effective = group_id.unwrap_or("default");
let rows = sqlx::query_as::<_, RawRow>(
"SELECT id, agent_id, source, tool_pattern, path_pattern, action, note, priority, group_id
FROM approval_rules
WHERE group_id = ?1 OR group_id = 'default'
ORDER BY CASE WHEN group_id = ?1 THEN 0 ELSE 1 END, priority ASC, id ASC",
)
.bind(effective)
.fetch_all(pool)
.await?;
rows.into_iter().map(from_raw).collect()
}
/// Inserts a new rule; returns its id.
pub async fn insert(pool: &SqlitePool, r: NewApprovalRule) -> Result<i64> {
let priority = r.priority.unwrap_or(100);
let group_id = r.group_id.as_deref().unwrap_or("default");
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO approval_rules (agent_id, source, tool_pattern, path_pattern, action, note, priority, group_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
RETURNING id",
)
.bind(r.agent_id)
.bind(r.source)
.bind(r.tool_pattern)
.bind(r.path_pattern)
.bind(r.action.as_str())
.bind(r.note)
.bind(priority)
.bind(group_id)
.fetch_one(pool)
.await?;
Ok(id)
}
/// Updates an existing rule by id.
pub async fn update(pool: &SqlitePool, id: i64, r: NewApprovalRule) -> Result<()> {
let priority = r.priority.unwrap_or(100);
let group_id = r.group_id.as_deref().unwrap_or("default");
sqlx::query(
"UPDATE approval_rules
SET agent_id = ?, source = ?, tool_pattern = ?, path_pattern = ?, action = ?, note = ?, priority = ?, group_id = ?
WHERE id = ?",
)
.bind(r.agent_id)
.bind(r.source)
.bind(r.tool_pattern)
.bind(r.path_pattern)
.bind(r.action.as_str())
.bind(r.note)
.bind(priority)
.bind(group_id)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
/// Deletes a rule by id.
pub async fn delete(pool: &SqlitePool, id: i64) -> Result<()> {
sqlx::query("DELETE FROM approval_rules WHERE id = ?")
.bind(id)
.execute(pool)
.await?;
Ok(())
}
+276
View File
@@ -0,0 +1,276 @@
use sqlx::SqlitePool;
use core_api::message_meta::MessageMetadata;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Role {
User,
Assistant,
/// Invocation message from a calling agent to a sub-agent; mapped to `user`
/// when rebuilding LLM context, invisible in the UI.
Agent,
}
impl Role {
pub fn as_str(&self) -> &'static str {
match self {
Role::User => "user",
Role::Assistant => "assistant",
Role::Agent => "agent",
}
}
pub fn from_str(s: &str) -> anyhow::Result<Self> {
match s {
"user" => Ok(Role::User),
"assistant" => Ok(Role::Assistant),
"agent" => Ok(Role::Agent),
other => anyhow::bail!("Unknown role: {other}"),
}
}
}
#[derive(Debug, Clone)]
pub struct ChatMessage {
pub id: i64,
pub role: Role,
pub content: String,
pub status: String,
pub input_tokens: Option<i64>,
pub output_tokens: Option<i64>,
/// True for messages injected synthetically (e.g. TIC notifications) — not
/// typed by a real user. Stored in DB so the UI can skip them on reload.
pub is_synthetic: bool,
/// Chain-of-thought from reasoning models (e.g. DeepSeek thinking mode).
/// Null for all other providers.
pub reasoning_content: Option<String>,
/// Cost of the turn in USD, when the provider reports it (OpenRouter).
/// Null for providers that don't bill per-request.
pub cost: Option<f64>,
/// Generic structured metadata (JSON column): file attachments today,
/// extensible later. `None` when the row has no metadata.
pub metadata: Option<MessageMetadata>,
pub created_at: Option<String>,
}
/// Raw row tuple for the shared `SELECT` projection. sqlx 0.9 requires SQL to be
/// `&'static str`, so the column list is repeated literally in each query below;
/// keep it in sync with this tuple and [`row_to_message`].
type Row = (
i64, String, String, String, Option<i64>, Option<i64>, bool,
Option<String>, Option<f64>, Option<String>, Option<String>,
);
/// Maps a [`Row`] into a [`ChatMessage`]. Metadata that fails to parse is treated
/// as absent (defensive: a malformed blob must not break history loading).
fn row_to_message(r: Row) -> anyhow::Result<ChatMessage> {
let (id, role, content, status, input_tokens, output_tokens, is_synthetic, reasoning_content, cost, metadata, created_at) = r;
Ok(ChatMessage {
id,
role: Role::from_str(&role)?,
content,
status,
input_tokens,
output_tokens,
is_synthetic,
reasoning_content,
cost,
metadata: metadata.and_then(|s| serde_json::from_str(&s).ok()),
created_at,
})
}
/// Appends a message with no structured metadata (the common case).
pub async fn append(
pool: &SqlitePool,
session_stack_id: i64,
role: &Role,
content: &str,
is_synthetic: bool,
reasoning_content: Option<&str>,
) -> anyhow::Result<i64> {
append_with_metadata(pool, session_stack_id, role, content, is_synthetic, reasoning_content, None).await
}
/// Like [`append`] but persists optional structured [`MessageMetadata`] (e.g. file
/// attachments) as a JSON blob. Empty metadata is stored as `NULL`.
pub async fn append_with_metadata(
pool: &SqlitePool,
session_stack_id: i64,
role: &Role,
content: &str,
is_synthetic: bool,
reasoning_content: Option<&str>,
metadata: Option<&MessageMetadata>,
) -> anyhow::Result<i64> {
let metadata_json = metadata
.filter(|m| !m.is_empty())
.map(serde_json::to_string)
.transpose()?;
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO chat_history (session_stack_id, role, content, is_synthetic, reasoning_content, metadata) \
VALUES (?, ?, ?, ?, ?, ?) RETURNING id",
)
.bind(session_stack_id)
.bind(role.as_str())
.bind(content)
.bind(is_synthetic as i64)
.bind(reasoning_content)
.bind(metadata_json)
.fetch_one(pool)
.await?;
Ok(id)
}
pub async fn mark_failed(pool: &SqlitePool, id: i64) -> anyhow::Result<()> {
sqlx::query("UPDATE chat_history SET status = 'failed' WHERE id = ?")
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn set_usage(
pool: &SqlitePool,
id: i64,
input_tokens: u32,
output_tokens: u32,
duration_ms: u64,
cost: Option<f64>,
) -> anyhow::Result<()> {
sqlx::query(
"UPDATE chat_history
SET input_tokens = ?, output_tokens = ?, duration_ms = ?, cost = ?
WHERE id = ?",
)
.bind(input_tokens as i64)
.bind(output_tokens as i64)
.bind(duration_ms as i64)
.bind(cost)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
/// All ok messages for a stack frame, ordered chronologically.
/// Used to rebuild LLM context for a specific agent.
pub async fn for_stack(
pool: &SqlitePool,
session_stack_id: i64,
) -> anyhow::Result<Vec<ChatMessage>> {
let rows = sqlx::query_as::<_, Row>(
"SELECT id, role, content, status, input_tokens, output_tokens, is_synthetic, reasoning_content, cost, metadata, created_at
FROM chat_history
WHERE session_stack_id = ? AND status = 'ok'
ORDER BY id ASC",
)
.bind(session_stack_id)
.fetch_all(pool)
.await?;
rows.into_iter().map(row_to_message).collect()
}
/// All messages for a stack frame including failed ones, ordered chronologically.
/// Used by the UI history API so the user can see cancelled messages.
pub async fn for_stack_all(
pool: &SqlitePool,
session_stack_id: i64,
) -> anyhow::Result<Vec<ChatMessage>> {
let rows = sqlx::query_as::<_, Row>(
"SELECT id, role, content, status, input_tokens, output_tokens, is_synthetic, reasoning_content, cost, metadata, created_at
FROM chat_history
WHERE session_stack_id = ?
ORDER BY id ASC",
)
.bind(session_stack_id)
.fetch_all(pool)
.await?;
rows.into_iter().map(row_to_message).collect()
}
/// Ok messages for a stack frame whose id is strictly greater than `after_id`,
/// ordered chronologically. Used by `build_openai_messages` when a compaction
/// summary exists: only the "raw" messages after the summary boundary are loaded.
pub async fn for_stack_since(
pool: &SqlitePool,
session_stack_id: i64,
after_id: i64,
) -> anyhow::Result<Vec<ChatMessage>> {
let rows = sqlx::query_as::<_, Row>(
"SELECT id, role, content, status, input_tokens, output_tokens, is_synthetic, reasoning_content, cost, metadata, created_at
FROM chat_history
WHERE session_stack_id = ? AND status = 'ok' AND id > ?
ORDER BY id ASC",
)
.bind(session_stack_id)
.bind(after_id)
.fetch_all(pool)
.await?;
rows.into_iter().map(row_to_message).collect()
}
/// Returns the most recent ok message for a stack frame, or `None` if empty.
/// Used by Telegram's `/context` command to show last turn's token usage.
pub async fn last_message_for_stack(
pool: &SqlitePool,
session_stack_id: i64,
) -> anyhow::Result<Option<ChatMessage>> {
let row = sqlx::query_as::<_, Row>(
"SELECT id, role, content, status, input_tokens, output_tokens, is_synthetic, reasoning_content, cost, metadata, created_at
FROM chat_history
WHERE session_stack_id = ? AND status = 'ok'
ORDER BY id DESC
LIMIT 1",
)
.bind(session_stack_id)
.fetch_optional(pool)
.await?;
row.map(row_to_message).transpose()
}
/// Total cost (USD) of a whole session: all messages across every stack frame
/// (main + sync sub-agents) that share this `session_id`. Async tasks live in
/// their own session and are naturally excluded. Returns `None` when no message
/// has a recorded cost (e.g. the provider does not report per-request pricing).
///
/// No `status` filter: money is spent even on turns later marked `failed`, so the
/// total reflects real spend. Uses plain `SUM(cost)` so an all-NULL set yields
/// `None`, distinguishing "no cost data" from a genuine `$0.00`.
pub async fn total_cost_for_session(
pool: &SqlitePool,
session_id: i64,
) -> anyhow::Result<Option<f64>> {
let total: Option<f64> = sqlx::query_scalar(
"SELECT SUM(ch.cost)
FROM chat_history ch
JOIN chat_sessions_stack css ON ch.session_stack_id = css.id
WHERE css.session_id = ?",
)
.bind(session_id)
.fetch_one(pool)
.await?;
Ok(total)
}
/// Rough token estimate for a stack frame (sum of content lengths / 4).
/// Used as a fallback when the LLM provider does not return usage data.
pub async fn estimate_tokens_for_stack(
pool: &SqlitePool,
session_stack_id: i64,
) -> anyhow::Result<u32> {
let total_chars: i64 = sqlx::query_scalar(
"SELECT COALESCE(SUM(LENGTH(content)), 0)
FROM chat_history
WHERE session_stack_id = ? AND status = 'ok'",
)
.bind(session_stack_id)
.fetch_one(pool)
.await?;
Ok((total_chars / 4).max(0) as u32)
}
+144
View File
@@ -0,0 +1,144 @@
use sqlx::SqlitePool;
#[derive(Debug, Clone)]
pub struct LlmToolCall {
pub id: i64,
pub message_id: i64,
pub name: String,
pub arguments: Option<String>,
pub result: Option<String>,
/// Result type tag: `"string"` (plain text, default) or `"json"` (structured
/// payload, e.g. MCP `structuredContent`). Drives frontend rendering.
pub result_type: String,
pub status: String,
}
/// Inserts a tool call in `running` state and returns its id.
/// `message_id` is the assistant `chat_history` row that triggered the call.
pub async fn append(
pool: &SqlitePool,
message_id: i64,
name: &str,
arguments: &str,
) -> anyhow::Result<i64> {
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO chat_llm_tools (message_id, name, arguments, status) VALUES (?, ?, ?, 'running') RETURNING id",
)
.bind(message_id)
.bind(name)
.bind(arguments)
.fetch_one(pool)
.await?;
Ok(id)
}
/// Marks a tool call as `pending` (waiting for explicit user approval or clarification).
/// Called just before registering an approval/clarification channel so `'pending'`
/// in the DB means "blocked on user input", not "still executing".
pub async fn set_approval_pending(pool: &SqlitePool, id: i64) -> anyhow::Result<()> {
sqlx::query("UPDATE chat_llm_tools SET status='pending' WHERE id=?")
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn complete(pool: &SqlitePool, id: i64, result: &str, result_type: &str) -> anyhow::Result<()> {
sqlx::query(
"UPDATE chat_llm_tools SET result = ?, result_type = ?, status = 'done' WHERE id = ?",
)
.bind(result)
.bind(result_type)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn fail(pool: &SqlitePool, id: i64, error: &str) -> anyhow::Result<()> {
sqlx::query(
"UPDATE chat_llm_tools SET result = ?, status = 'failed' WHERE id = ?",
)
.bind(error)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
/// Marks a tool call as `cancelled` — stopped by the user via `/stop`.
/// Terminal and distinct from `failed`: a cancellation is deliberate, not an
/// error, and is **not** picked up by `pending_for_stack` (never re-run on
/// restart, unlike an interrupted `running` call).
pub async fn cancel(pool: &SqlitePool, id: i64, note: &str) -> anyhow::Result<()> {
sqlx::query(
"UPDATE chat_llm_tools SET result = ?, status = 'cancelled' WHERE id = ?",
)
.bind(note)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
/// Marks a tool call as `rejected` — denied by an approval policy or a human.
/// Terminal and distinct from `failed`: a denial is a policy decision, not an
/// error, and is not re-run on restart.
pub async fn reject(pool: &SqlitePool, id: i64, reason: &str) -> anyhow::Result<()> {
sqlx::query(
"UPDATE chat_llm_tools SET result = ?, status = 'rejected' WHERE id = ?",
)
.bind(reason)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
/// All `running` or `pending` tool calls for a stack frame — used to resume interrupted sessions.
/// `running`: tool was executing when the session was interrupted (re-execute).
/// `pending`: tool was waiting for explicit user approval or clarification (re-gate or re-ask).
pub async fn pending_for_stack(
pool: &SqlitePool,
session_stack_id: i64,
) -> anyhow::Result<Vec<LlmToolCall>> {
let rows = sqlx::query_as::<_, (i64, i64, String, Option<String>, Option<String>, String, String)>(
"SELECT t.id, t.message_id, t.name, t.arguments, t.result, t.result_type, t.status
FROM chat_llm_tools t
JOIN chat_history h ON t.message_id = h.id
WHERE h.session_stack_id = ?
AND t.status IN ('running', 'pending')
ORDER BY t.id ASC",
)
.bind(session_stack_id)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(row_to_tool).collect())
}
/// All tool calls for a single assistant message, ordered chronologically.
pub async fn for_message(
pool: &SqlitePool,
message_id: i64,
) -> anyhow::Result<Vec<LlmToolCall>> {
let rows = sqlx::query_as::<_, (i64, i64, String, Option<String>, Option<String>, String, String)>(
"SELECT id, message_id, name, arguments, result, result_type, status
FROM chat_llm_tools
WHERE message_id = ?
ORDER BY id ASC",
)
.bind(message_id)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(row_to_tool).collect())
}
fn row_to_tool(
(id, message_id, name, arguments, result, result_type, status): (
i64, i64, String, Option<String>, Option<String>, String, String,
),
) -> LlmToolCall {
LlmToolCall { id, message_id, name, arguments, result, result_type, status }
}
+76
View File
@@ -0,0 +1,76 @@
use sqlx::SqlitePool;
pub struct ChatSession {
pub id: i64,
pub source: String,
pub agent_id: String,
/// True when a real user is actively participating (web, telegram).
/// False for fully automated sessions (cron, tic).
pub is_interactive: bool,
/// True for short-lived task sessions (cron, tic) with no long-term
/// conversational value. May be used to skip memory / analytics sinks.
pub is_ephemeral: bool,
/// Optional RunContext JSON blob assigned to this session.
/// `None` resolves to the implicit "default" run_context at runtime.
pub run_context: Option<String>,
}
pub async fn create(
pool: &SqlitePool,
agent_id: &str,
source: &str,
is_interactive: bool,
is_ephemeral: bool,
) -> anyhow::Result<ChatSession> {
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO chat_sessions (source, agent_id, is_interactive, is_ephemeral)
VALUES (?, ?, ?, ?) RETURNING id",
)
.bind(source)
.bind(agent_id)
.bind(is_interactive as i64)
.bind(is_ephemeral as i64)
.fetch_one(pool)
.await?;
Ok(ChatSession {
id,
source: source.to_string(),
agent_id: agent_id.to_string(),
is_interactive,
is_ephemeral,
run_context: None,
})
}
pub async fn set_run_context(
pool: &SqlitePool,
id: i64,
run_context: Option<&str>,
) -> anyhow::Result<()> {
sqlx::query("UPDATE chat_sessions SET run_context = ? WHERE id = ?")
.bind(run_context)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn find_by_id(pool: &SqlitePool, id: i64) -> anyhow::Result<Option<ChatSession>> {
let row = sqlx::query_as::<_, (i64, String, String, bool, bool, Option<String>)>(
"SELECT id, source, agent_id, is_interactive, is_ephemeral, run_context
FROM chat_sessions WHERE id = ?",
)
.bind(id)
.fetch_optional(pool)
.await?;
Ok(row.map(|(id, source, agent_id, is_interactive, is_ephemeral, run_context)| ChatSession {
id,
source,
agent_id,
is_interactive,
is_ephemeral,
run_context,
}))
}
@@ -0,0 +1,190 @@
use sqlx::SqlitePool;
#[derive(Debug, Clone)]
pub struct SessionStack {
pub id: i64,
pub agent_id: String,
pub depth: i64,
pub parent_tool_call_id: Option<i64>,
}
pub async fn create(
pool: &SqlitePool,
session_id: i64,
agent_id: &str,
agent_prompt: Option<&str>,
depth: i64,
parent_tool_call_id: Option<i64>,
) -> anyhow::Result<SessionStack> {
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO chat_sessions_stack (session_id, agent_id, agent_prompt, depth, parent_tool_call_id)
VALUES (?, ?, ?, ?, ?) RETURNING id",
)
.bind(session_id)
.bind(agent_id)
.bind(agent_prompt)
.bind(depth)
.bind(parent_tool_call_id)
.fetch_one(pool)
.await?;
Ok(SessionStack { id, agent_id: agent_id.to_string(), depth, parent_tool_call_id })
}
/// Returns the deepest active (non-terminated) frame for a session.
pub async fn active_for_session(
pool: &SqlitePool,
session_id: i64,
) -> anyhow::Result<Option<SessionStack>> {
let row = sqlx::query_as::<_, (i64, String, i64, Option<i64>)>(
"SELECT id, agent_id, depth, parent_tool_call_id
FROM chat_sessions_stack
WHERE session_id = ?
AND terminated_at IS NULL
ORDER BY depth DESC
LIMIT 1",
)
.bind(session_id)
.fetch_optional(pool)
.await?;
Ok(row.map(row_to_stack))
}
/// All active (non-terminated) frames for a session, ordered by depth ASC.
/// Used by restart recovery to detect an interrupted parallel sub-agent batch:
/// a purely linear stack has at most one active frame per depth, so ≥2 active
/// frames at the same depth can only be a concurrent batch left mid-flight.
pub async fn active_all_for_session(
pool: &SqlitePool,
session_id: i64,
) -> anyhow::Result<Vec<SessionStack>> {
let rows = sqlx::query_as::<_, (i64, String, i64, Option<i64>)>(
"SELECT id, agent_id, depth, parent_tool_call_id
FROM chat_sessions_stack
WHERE session_id = ?
AND terminated_at IS NULL
ORDER BY depth ASC",
)
.bind(session_id)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(row_to_stack).collect())
}
/// Returns the root (depth=0) stack frame for a session.
pub async fn main_for_session(
pool: &SqlitePool,
session_id: i64,
) -> anyhow::Result<Option<SessionStack>> {
let row = sqlx::query_as::<_, (i64, String, i64, Option<i64>)>(
"SELECT id, agent_id, depth, parent_tool_call_id
FROM chat_sessions_stack
WHERE session_id = ? AND depth = 0
ORDER BY id ASC
LIMIT 1",
)
.bind(session_id)
.fetch_optional(pool)
.await?;
Ok(row.map(row_to_stack))
}
/// Returns all stack frames for a session (including terminated), ordered by id ASC.
/// Used to reconstruct the full agent call tree from history.
pub async fn all_for_session(
pool: &SqlitePool,
session_id: i64,
) -> anyhow::Result<Vec<SessionStack>> {
let rows = sqlx::query_as::<_, (i64, String, i64, Option<i64>)>(
"SELECT id, agent_id, depth, parent_tool_call_id
FROM chat_sessions_stack
WHERE session_id = ?
ORDER BY id ASC",
)
.bind(session_id)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(row_to_stack).collect())
}
pub async fn find_by_id(pool: &SqlitePool, id: i64) -> anyhow::Result<Option<SessionStack>> {
let row = sqlx::query_as::<_, (i64, String, i64, Option<i64>)>(
"SELECT id, agent_id, depth, parent_tool_call_id
FROM chat_sessions_stack
WHERE id = ?",
)
.bind(id)
.fetch_optional(pool)
.await?;
Ok(row.map(row_to_stack))
}
/// Marks a stack frame as terminated (agent completed or was cancelled).
pub async fn terminate(pool: &SqlitePool, id: i64) -> anyhow::Result<()> {
sqlx::query(
"UPDATE chat_sessions_stack SET terminated_at = datetime('now') WHERE id = ?",
)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
fn row_to_stack(
(id, agent_id, depth, parent_tool_call_id): (i64, String, i64, Option<i64>),
) -> SessionStack {
SessionStack { id, agent_id, depth, parent_tool_call_id }
}
#[cfg(test)]
mod tests {
use super::*;
fn temp_db_path(tag: &str) -> String {
let mut p = std::env::temp_dir();
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
p.push(format!("skald-test-{tag}-{}-{nanos}.db", std::process::id()));
p.to_string_lossy().into_owned()
}
fn cleanup(path: &str) {
for suffix in ["", "-wal", "-shm"] {
let _ = std::fs::remove_file(format!("{path}{suffix}"));
}
}
/// Validates `active_all_for_session` against the real schema: two active
/// frames at the same depth are the signature of an interrupted parallel
/// sub-agent batch, and terminating one drops it from the active set.
#[tokio::test]
async fn active_all_reflects_parallel_siblings() {
let path = temp_db_path("stack-parallel");
let pool = crate::db::init_system_pool(&path).await.unwrap();
let sid = 1;
// `session_id` has a FK to chat_sessions (sqlx enables foreign_keys).
sqlx::query("INSERT INTO chat_sessions (id) VALUES (?)")
.bind(sid).execute(&pool).await.unwrap();
create(&pool, sid, "main", None, 0, None).await.unwrap();
let a = create(&pool, sid, "task", Some("A"), 1, Some(101)).await.unwrap();
create(&pool, sid, "task", Some("B"), 1, Some(102)).await.unwrap();
let active = active_all_for_session(&pool, sid).await.unwrap();
assert_eq!(active.len(), 3, "root + two live siblings");
assert_eq!(active.iter().filter(|f| f.depth == 1).count(), 2, "two frames share depth 1");
// Terminating one sibling removes it from the active set (used by reap).
terminate(&pool, a.id).await.unwrap();
let active = active_all_for_session(&pool, sid).await.unwrap();
assert_eq!(active.len(), 2);
assert!(active.iter().all(|f| f.id != a.id), "terminated frame is excluded");
pool.close().await;
cleanup(&path);
}
}
@@ -0,0 +1,64 @@
//! Persistent conversation summaries generated by the context compactor.
//!
//! Each row covers all `chat_history` messages up to and including
//! `covers_up_to_message_id`. At most one active summary exists per stack;
//! a new compaction creates a new row that supersedes the previous one.
use sqlx::SqlitePool;
#[derive(Debug, Clone)]
pub struct ChatSummary {
pub id: i64,
pub stack_id: i64,
pub content: String,
/// All chat_history rows with `id <= covers_up_to_message_id` are covered
/// by this summary. `build_openai_messages` loads only rows *after* this id.
pub covers_up_to_message_id: i64,
pub created_at: String,
}
/// Persist a new summary for the given stack, returning the new row id.
pub async fn save(
pool: &SqlitePool,
stack_id: i64,
content: &str,
covers_up_to_message_id: i64,
) -> anyhow::Result<i64> {
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO chat_summaries
(stack_id, content, covers_up_to_message_id)
VALUES (?, ?, ?)
RETURNING id",
)
.bind(stack_id)
.bind(content)
.bind(covers_up_to_message_id)
.fetch_one(pool)
.await?;
Ok(id)
}
/// Returns the most recent summary for a stack, or `None` if none exists.
pub async fn latest_for_stack(
pool: &SqlitePool,
stack_id: i64,
) -> anyhow::Result<Option<ChatSummary>> {
let row = sqlx::query_as::<_, (i64, i64, String, i64, String)>(
"SELECT id, stack_id, content, covers_up_to_message_id, created_at
FROM chat_summaries
WHERE stack_id = ?
ORDER BY id DESC
LIMIT 1",
)
.bind(stack_id)
.fetch_optional(pool)
.await?;
Ok(row.map(|(id, stack_id, content, covers_up_to_message_id, created_at)| ChatSummary {
id,
stack_id,
content,
covers_up_to_message_id,
created_at,
}))
}
+44
View File
@@ -0,0 +1,44 @@
use sqlx::SqlitePool;
pub struct ConfigEntry {
pub key: String,
pub value: String,
pub updated_at: String,
}
/// Get a config value by key.
pub async fn get(pool: &SqlitePool, key: &str) -> anyhow::Result<Option<String>> {
let row = sqlx::query_as::<_, (String,)>(
"SELECT value FROM config WHERE key = ?",
)
.bind(key)
.fetch_optional(pool)
.await?;
Ok(row.map(|(v,)| v))
}
/// Upsert a config key/value pair.
pub async fn set(pool: &SqlitePool, key: &str, value: &str) -> anyhow::Result<()> {
sqlx::query(
"INSERT INTO config (key, value, updated_at)
VALUES (?, ?, datetime('now'))
ON CONFLICT(key) DO UPDATE SET
value = excluded.value,
updated_at = excluded.updated_at",
)
.bind(key)
.bind(value)
.execute(pool)
.await?;
Ok(())
}
/// Delete a config entry.
pub async fn delete(pool: &SqlitePool, key: &str) -> anyhow::Result<()> {
sqlx::query("DELETE FROM config WHERE key = ?")
.bind(key)
.execute(pool)
.await?;
Ok(())
}
+103
View File
@@ -0,0 +1,103 @@
use anyhow::Result;
use sqlx::SqlitePool;
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct JobRun {
pub id: i64,
pub job_id: i64,
pub session_id: Option<i64>,
pub started_at: String,
pub completed_at: Option<String>,
pub duration_ms: Option<i64>,
pub status: String,
pub final_response: Option<String>,
pub error: Option<String>,
pub created_at: String,
}
pub async fn insert(
pool: &SqlitePool,
job_id: i64,
session_id: Option<i64>,
started_at: &str,
completed_at: &str,
duration_ms: i64,
status: &str,
final_response: Option<&str>,
error: Option<&str>,
) -> Result<JobRun> {
let id = sqlx::query(
"INSERT INTO job_runs (job_id, session_id, started_at, completed_at, duration_ms, status, final_response, error)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind(job_id)
.bind(session_id)
.bind(started_at)
.bind(completed_at)
.bind(duration_ms)
.bind(status)
.bind(final_response)
.bind(error)
.execute(pool)
.await?
.last_insert_rowid();
let row = sqlx::query_as::<_, JobRun>(
"SELECT id, job_id, session_id, started_at, completed_at, duration_ms,
status, final_response, error, created_at
FROM job_runs WHERE id = ?",
)
.bind(id)
.fetch_one(pool)
.await?;
Ok(row)
}
pub async fn list_for_job(pool: &SqlitePool, job_id: i64, limit: i64) -> Result<Vec<JobRun>> {
let rows = sqlx::query_as::<_, JobRun>(
"SELECT id, job_id, session_id, started_at, completed_at, duration_ms,
status, final_response, error, created_at
FROM job_runs
WHERE job_id = ?
ORDER BY created_at DESC
LIMIT ?",
)
.bind(job_id)
.bind(limit)
.fetch_all(pool)
.await?;
Ok(rows)
}
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct JobRunWithMeta {
pub id: i64,
pub job_id: i64,
pub session_id: Option<i64>,
pub started_at: String,
pub completed_at: Option<String>,
pub duration_ms: Option<i64>,
pub status: String,
pub final_response: Option<String>,
pub error: Option<String>,
pub created_at: String,
pub job_title: Option<String>,
pub agent_id: Option<String>,
pub kind: Option<String>,
}
pub async fn list_all(pool: &SqlitePool, limit: i64) -> Result<Vec<JobRunWithMeta>> {
let rows = sqlx::query_as::<_, JobRunWithMeta>(
"SELECT jr.id, jr.job_id, jr.session_id, jr.started_at, jr.completed_at,
jr.duration_ms, jr.status, jr.final_response, jr.error, jr.created_at,
sj.title AS job_title, sj.agent_id, sj.kind
FROM job_runs jr
LEFT JOIN scheduled_jobs sj ON jr.job_id = sj.id
ORDER BY jr.created_at DESC
LIMIT ?",
)
.bind(limit)
.fetch_all(pool)
.await?;
Ok(rows)
}
+94
View File
@@ -0,0 +1,94 @@
//! `known_tools` — every tool ever offered to the LLM, captured at injection
//! time by [`crate::tool_discovery::ToolDiscovery`].
//!
//! This is the drift-proof half of tool visibility: instead of maintaining a
//! parallel list of "all tools", we record what is actually assembled into the
//! LLM request (`AgentRunConfig::all_tool_defs`). The approval / Security-groups
//! UI merges these rows so tools injected outside the `ToolRegistry` (interface
//! tools, plugin tools, provider tools) can still be assigned a permission.
use anyhow::Result;
use serde::Serialize;
use sqlx::SqlitePool;
#[derive(Debug, Clone, Serialize)]
pub struct KnownTool {
pub name: String,
pub description: String,
/// JSON parameters schema as last seen, if any.
pub schema: Option<String>,
}
/// Records (or refreshes) a tool by name. Idempotent: re-seeing a tool updates
/// its description/schema and bumps `last_seen`.
pub async fn upsert(
pool: &SqlitePool,
name: &str,
description: &str,
schema: Option<&str>,
) -> Result<()> {
sqlx::query(
"INSERT INTO known_tools (name, description, schema, first_seen, last_seen)
VALUES (?1, ?2, ?3, strftime('%s','now'), strftime('%s','now'))
ON CONFLICT(name) DO UPDATE SET
description = excluded.description,
schema = excluded.schema,
last_seen = excluded.last_seen",
)
.bind(name)
.bind(description)
.bind(schema)
.execute(pool)
.await?;
Ok(())
}
/// All recorded tools, sorted by name.
pub async fn all(pool: &SqlitePool) -> Result<Vec<KnownTool>> {
let rows = sqlx::query_as::<_, (String, String, Option<String>)>(
"SELECT name, description, schema FROM known_tools ORDER BY name",
)
.fetch_all(pool)
.await?;
Ok(rows
.into_iter()
.map(|(name, description, schema)| KnownTool { name, description, schema })
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
fn temp_db_path(tag: &str) -> String {
let mut p = std::env::temp_dir();
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
p.push(format!("skald-test-{tag}-{}-{nanos}.db", std::process::id()));
p.to_string_lossy().into_owned()
}
fn cleanup(path: &str) {
for suffix in ["", "-wal", "-shm"] {
let _ = std::fs::remove_file(format!("{path}{suffix}"));
}
}
#[tokio::test]
async fn upsert_is_idempotent_and_updates_metadata() {
let path = temp_db_path("known-tools");
let pool = crate::db::init_system_pool(&path).await.unwrap();
upsert(&pool, "send_voice_message", "v1", Some(r#"{"type":"object"}"#)).await.unwrap();
upsert(&pool, "send_voice_message", "v2", None).await.unwrap();
let rows = all(&pool).await.unwrap();
assert_eq!(rows.len(), 1, "same name must not create a second row");
assert_eq!(rows[0].name, "send_voice_message");
assert_eq!(rows[0].description, "v2", "description is refreshed on re-upsert");
assert_eq!(rows[0].schema, None, "schema is refreshed on re-upsert");
pool.close().await;
cleanup(&path);
}
}
@@ -0,0 +1,73 @@
//! Background maintenance task for the `llm_requests` table.
//!
//! Periodically nulls out old payloads/headers and deletes expired rows according
//! to the retention settings in [`LlmRequestsLogConfig`], then `VACUUM`s to reclaim
//! freed pages. Extracted from `Skald::new` so the loop lives next to the queries it
//! calls; the returned handle is registered with the `TaskSupervisor` for shutdown.
use std::sync::Arc;
use std::time::Duration;
use sqlx::SqlitePool;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use crate::config::LlmRequestsLogConfig;
/// Spawns the retention/cleanup loop for the `llm_requests` table.
///
/// First run happens 1 minute after startup, then every 12 hours. The loop exits
/// when `shutdown` is cancelled. Callers should register the returned handle with
/// the task supervisor so it is awaited on shutdown.
pub fn spawn(
pool: Arc<SqlitePool>,
cfg: LlmRequestsLogConfig,
shutdown: CancellationToken,
) -> JoinHandle<()> {
tokio::spawn(async move {
tokio::select! {
_ = shutdown.cancelled() => { return; }
_ = tokio::time::sleep(Duration::from_secs(60)) => {}
}
loop {
if let Some(days) = cfg.cleanup_request_payload_after {
match super::null_request_payload(&pool, days).await {
Ok(n) if n > 0 => info!(rows = n, days, "llm_requests: nulled request payload"),
Ok(_) => {}
Err(e) => warn!(error = %e, "llm_requests: null request payload failed"),
}
}
if let Some(days) = cfg.cleanup_response_payload_after {
match super::null_response_payload(&pool, days).await {
Ok(n) if n > 0 => info!(rows = n, days, "llm_requests: nulled response payload"),
Ok(_) => {}
Err(e) => warn!(error = %e, "llm_requests: null response payload failed"),
}
}
if let Some(days) = cfg.cleanup_headers_after {
match super::null_headers(&pool, days).await {
Ok(n) if n > 0 => info!(rows = n, days, "llm_requests: nulled headers"),
Ok(_) => {}
Err(e) => warn!(error = %e, "llm_requests: null headers failed"),
}
}
if let Some(days) = cfg.cleanup_rows_after {
match super::delete_old_rows(&pool, days).await {
Ok(n) if n > 0 => info!(deleted = n, days, "llm_requests: deleted old rows"),
Ok(_) => {}
Err(e) => warn!(error = %e, "llm_requests: delete old rows failed"),
}
}
// VACUUM reclaims pages freed by DELETE/UPDATE NULL.
match sqlx::query("VACUUM").execute(&*pool).await {
Ok(_) => info!("llm_requests: VACUUM complete"),
Err(e) => warn!(error = %e, "llm_requests: VACUUM failed"),
}
tokio::select! {
_ = shutdown.cancelled() => { break; }
_ = tokio::time::sleep(Duration::from_secs(12 * 3600)) => {}
}
}
})
}
@@ -0,0 +1,125 @@
//! DB operations for the `llm_requests` table.
//!
//! Every `chat_with_tools` call is logged here by the
//! [`crate::chatbot::logging::LoggingChatbotClient`] wrapper.
//! Rows are retained for `llm.request_log.retention_days` days (default 14).
use anyhow::Result;
use sqlx::SqlitePool;
pub mod cleanup;
// ── Row struct ────────────────────────────────────────────────────────────────
pub struct LlmRequestRow {
pub session_id: Option<i64>,
pub stack_id: Option<i64>,
pub model_name: String,
/// Full HTTP request body sent to the provider (compact JSON, no pretty-print).
pub request_json: String,
/// HTTP request headers as a compact JSON object (api-key redacted).
pub request_headers: Option<String>,
/// Full HTTP response body from the provider (compact JSON).
pub response_json: Option<String>,
/// HTTP response headers as a compact JSON object.
pub response_headers: Option<String>,
/// Error message when the HTTP call itself failed (no response available).
pub error_text: Option<String>,
pub input_tokens: Option<i64>,
pub output_tokens: Option<i64>,
/// Wall-clock time of the full HTTP round-trip in milliseconds.
pub duration_ms: i64,
/// Tokens served from the provider's prompt cache (already parsed by the client).
pub cache_read_tokens: Option<i64>,
/// Tokens written into the provider's prompt cache (Anthropic only).
pub cache_creation_tokens: Option<i64>,
}
// ── Writes ────────────────────────────────────────────────────────────────────
pub async fn insert(pool: &SqlitePool, row: LlmRequestRow) -> Result<i64> {
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO llm_requests (
session_id, stack_id, model_name,
request_json, request_headers,
response_json, response_headers,
error_text, input_tokens, output_tokens, duration_ms,
cache_read_tokens, cache_creation_tokens
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
RETURNING id",
)
.bind(row.session_id)
.bind(row.stack_id)
.bind(&row.model_name)
.bind(&row.request_json)
.bind(&row.request_headers)
.bind(&row.response_json)
.bind(&row.response_headers)
.bind(&row.error_text)
.bind(row.input_tokens)
.bind(row.output_tokens)
.bind(row.duration_ms)
.bind(row.cache_read_tokens)
.bind(row.cache_creation_tokens)
.fetch_one(pool)
.await?;
Ok(id)
}
// ── Maintenance ───────────────────────────────────────────────────────────────
/// Physically deletes rows older than `days` days. Returns rows affected.
pub async fn delete_old_rows(pool: &SqlitePool, days: u32) -> Result<u64> {
let cutoff = format!("-{days} days");
let n = sqlx::query("DELETE FROM llm_requests WHERE created_at < datetime('now', ?)")
.bind(&cutoff)
.execute(pool)
.await?
.rows_affected();
Ok(n)
}
/// Nulls out `request_json` for rows older than `days` days. Returns rows affected.
pub async fn null_request_payload(pool: &SqlitePool, days: u32) -> Result<u64> {
let cutoff = format!("-{days} days");
let n = sqlx::query(
"UPDATE llm_requests SET request_json = '' \
WHERE request_json != '' AND created_at < datetime('now', ?)",
)
.bind(&cutoff)
.execute(pool)
.await?
.rows_affected();
Ok(n)
}
/// Nulls out `response_json` for rows older than `days` days. Returns rows affected.
pub async fn null_response_payload(pool: &SqlitePool, days: u32) -> Result<u64> {
let cutoff = format!("-{days} days");
let n = sqlx::query(
"UPDATE llm_requests SET response_json = NULL \
WHERE response_json IS NOT NULL AND created_at < datetime('now', ?)",
)
.bind(&cutoff)
.execute(pool)
.await?
.rows_affected();
Ok(n)
}
/// Nulls out both header columns for rows older than `days` days. Returns rows affected.
pub async fn null_headers(pool: &SqlitePool, days: u32) -> Result<u64> {
let cutoff = format!("-{days} days");
let n = sqlx::query(
"UPDATE llm_requests \
SET request_headers = NULL, response_headers = NULL \
WHERE (request_headers IS NOT NULL OR response_headers IS NOT NULL) \
AND created_at < datetime('now', ?)",
)
.bind(&cutoff)
.execute(pool)
.await?
.rows_affected();
Ok(n)
}
+114
View File
@@ -0,0 +1,114 @@
use anyhow::Result;
use sqlx::SqlitePool;
// ── Row type ──────────────────────────────────────────────────────────────────
pub struct McpEvent {
pub id: i64,
pub source: String,
pub method: String,
pub payload: String, // raw JSON of the "params" field
pub processed: bool,
pub processed_at: Option<String>,
pub created_at: String,
}
// ── Write ─────────────────────────────────────────────────────────────────────
/// Insert a new event (processed = false).
pub async fn insert(
pool: &SqlitePool,
source: &str,
method: &str,
payload: &str,
) -> Result<i64> {
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO mcp_events (source, method, payload)
VALUES (?, ?, ?)
RETURNING id",
)
.bind(source)
.bind(method)
.bind(payload)
.fetch_one(pool)
.await?;
Ok(id)
}
/// Mark a batch of events as processed (sets processed = 1, processed_at = now).
pub async fn mark_processed(pool: &SqlitePool, ids: &[i64]) -> Result<()> {
if ids.is_empty() {
return Ok(());
}
// Build a parameterised IN clause.
let placeholders = ids.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
let sql = format!(
"UPDATE mcp_events
SET processed = 1, processed_at = datetime('now')
WHERE id IN ({placeholders})"
);
let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
for id in ids {
q = q.bind(id);
}
q.execute(pool).await?;
Ok(())
}
// ── Read ──────────────────────────────────────────────────────────────────────
/// Oldest N pending (unprocessed) events, ordered oldest-first.
/// Used by TicManager to fetch a bounded batch each tick.
pub async fn pending_limited(pool: &SqlitePool, limit: i64) -> Result<Vec<McpEvent>> {
let rows = sqlx::query_as::<_, (i64, String, String, String, bool, Option<String>, String)>(
"SELECT id, source, method, payload, processed, processed_at, created_at
FROM mcp_events
WHERE processed = 0
ORDER BY created_at ASC
LIMIT ?",
)
.bind(limit)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(row_to_event).collect())
}
/// All pending (unprocessed) events, ordered oldest-first.
pub async fn pending(pool: &SqlitePool) -> Result<Vec<McpEvent>> {
let rows = sqlx::query_as::<_, (i64, String, String, String, bool, Option<String>, String)>(
"SELECT id, source, method, payload, processed, processed_at, created_at
FROM mcp_events
WHERE processed = 0
ORDER BY created_at ASC",
)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(row_to_event).collect())
}
/// All events (both processed and pending), most-recent first. Useful for debug/audit.
pub async fn all_recent(pool: &SqlitePool, limit: i64) -> Result<Vec<McpEvent>> {
let rows = sqlx::query_as::<_, (i64, String, String, String, bool, Option<String>, String)>(
"SELECT id, source, method, payload, processed, processed_at, created_at
FROM mcp_events
ORDER BY created_at DESC
LIMIT ?",
)
.bind(limit)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(row_to_event).collect())
}
// ── Helpers ───────────────────────────────────────────────────────────────────
fn row_to_event(
(id, source, method, payload, processed, processed_at, created_at): (
i64, String, String, String, bool, Option<String>, String,
),
) -> McpEvent {
McpEvent { id, source, method, payload, processed, processed_at, created_at }
}
+129
View File
@@ -0,0 +1,129 @@
use std::collections::HashMap;
use anyhow::Result;
use serde::{Deserialize, Serialize};
use sqlx::SqlitePool;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpServerRow {
pub id: i64,
pub name: String,
pub transport: String,
pub command: Option<String>,
pub args_json: Option<String>,
pub env_json: Option<String>,
pub url: Option<String>,
pub api_key: Option<String>,
pub description: Option<String>,
pub friendly_name: Option<String>,
pub enabled: bool,
}
impl McpServerRow {
pub fn args(&self) -> Vec<String> {
self.args_json.as_deref()
.and_then(|s| serde_json::from_str(s).ok())
.unwrap_or_default()
}
pub fn env(&self) -> HashMap<String, String> {
self.env_json.as_deref()
.and_then(|s| serde_json::from_str(s).ok())
.unwrap_or_default()
}
}
type RawRow = (i64, String, String, Option<String>, Option<String>, Option<String>, Option<String>, Option<String>, Option<String>, Option<String>, i64);
fn from_raw(r: RawRow) -> McpServerRow {
McpServerRow {
id: r.0,
name: r.1,
transport: r.2,
command: r.3,
args_json: r.4,
env_json: r.5,
url: r.6,
api_key: r.7,
description: r.8,
friendly_name: r.9,
enabled: r.10 != 0,
}
}
const SELECT: &str =
"SELECT id, name, transport, command, args_json, env_json, url, api_key, description, friendly_name, enabled \
FROM mcp_servers";
pub async fn all(pool: &SqlitePool) -> Result<Vec<McpServerRow>> {
let rows = sqlx::query_as::<_, RawRow>(sqlx::AssertSqlSafe(format!("{SELECT} ORDER BY name")))
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(from_raw).collect())
}
pub async fn all_enabled(pool: &SqlitePool) -> Result<Vec<McpServerRow>> {
let rows = sqlx::query_as::<_, RawRow>(sqlx::AssertSqlSafe(format!("{SELECT} WHERE enabled = 1 ORDER BY name")))
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(from_raw).collect())
}
pub struct UpsertParams<'a> {
pub name: &'a str,
pub transport: &'a str,
pub command: Option<&'a str>,
pub args_json: Option<String>,
pub env_json: Option<String>,
pub url: Option<&'a str>,
pub api_key: Option<&'a str>,
pub description: Option<&'a str>,
pub friendly_name: Option<&'a str>,
}
pub async fn upsert(pool: &SqlitePool, p: UpsertParams<'_>) -> Result<i64> {
let row = sqlx::query_as::<_, (i64,)>(
"INSERT INTO mcp_servers (name, transport, command, args_json, env_json, url, api_key, description, friendly_name, enabled)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, 1)
ON CONFLICT(name) DO UPDATE SET
transport = excluded.transport,
command = excluded.command,
args_json = excluded.args_json,
env_json = excluded.env_json,
url = excluded.url,
api_key = excluded.api_key,
description = excluded.description,
friendly_name = excluded.friendly_name,
enabled = 1
RETURNING id",
)
.bind(p.name)
.bind(p.transport)
.bind(p.command)
.bind(p.args_json)
.bind(p.env_json)
.bind(p.url)
.bind(p.api_key)
.bind(p.description)
.bind(p.friendly_name)
.fetch_one(pool)
.await?;
Ok(row.0)
}
pub async fn set_enabled(pool: &SqlitePool, name: &str, enabled: bool) -> Result<()> {
sqlx::query("UPDATE mcp_servers SET enabled = ?1 WHERE name = ?2")
.bind(enabled as i64)
.bind(name)
.execute(pool)
.await?;
Ok(())
}
pub async fn delete(pool: &SqlitePool, name: &str) -> Result<()> {
sqlx::query("DELETE FROM mcp_servers WHERE name = ?1")
.bind(name)
.execute(pool)
.await?;
Ok(())
}
+790
View File
@@ -0,0 +1,790 @@
pub mod approval_rules;
pub mod project_tickets;
pub mod projects;
pub mod chat_history;
pub mod chat_llm_tools;
pub mod chat_sessions;
pub mod chat_sessions_stack;
pub mod chat_summaries;
pub mod config;
pub mod job_runs;
pub mod known_tools;
pub mod llm_requests;
pub mod mcp_events;
pub mod mcp_servers;
pub mod plugins;
pub mod scheduled_jobs;
pub mod scratchpad;
pub mod session_mcp_grants;
pub mod sources;
pub mod stack_mcp_grants;
pub mod tool_permission_groups;
pub mod users;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::time::Duration;
use anyhow::{Context, Result};
use sqlx::{SqlitePool, sqlite::{SqliteConnectOptions, SqliteJournalMode, SqliteSynchronous}};
use crate::crypto::Dek;
/// Every database file lives here, so backup, export and per-user erasure are a
/// matter of files rather than tables.
pub const DATABASE_DIR: &str = "database";
/// System database: instance-wide state, shared by every user. Fixed path — not
/// configurable.
pub const SYSTEM_DB_PATH: &str = "database/system.db";
/// `{dir}/{userid}.db`. Keyed by the opaque user id and never the username, so a
/// rename never has to touch the file. `dir` is a parameter rather than the
/// [`DATABASE_DIR`] constant so nothing depends on the process's working
/// directory — tests in particular.
pub fn user_db_path(dir: &Path, user_id: &str) -> PathBuf {
dir.join(format!("{user_id}.db"))
}
/// The `-wal` and `-shm` sidecars SQLite keeps beside a database in WAL mode.
/// Erasing a user means erasing these too.
pub fn user_db_sidecars(path: &Path) -> [PathBuf; 2] {
let ext = |suffix: &str| {
let mut p = path.to_path_buf().into_os_string();
p.push(suffix);
PathBuf::from(p)
};
[ext("-wal"), ext("-shm")]
}
fn ensure_parent(path: &Path) -> Result<()> {
// `create_if_missing` creates the file, never its parent directory.
if let Some(parent) = path.parent()
&& !parent.as_os_str().is_empty()
{
std::fs::create_dir_all(parent)
.with_context(|| format!("failed to create database directory {}", parent.display()))?;
}
Ok(())
}
fn tuned(opts: SqliteConnectOptions) -> SqliteConnectOptions {
// WAL lets readers run alongside a single writer, and `busy_timeout` makes a
// writer *wait* for the lock instead of failing immediately with SQLITE_BUSY
// ("database is locked"). Without these, concurrent writers — e.g. the
// mobile-connector persisting its E2E `send_counter` while the chat loop /
// cron write history — abort mid-operation, which silently drops outbound
// mobile messages (inbox_update never reaches the device).
opts.journal_mode(SqliteJournalMode::Wal)
.synchronous(SqliteSynchronous::Normal)
.busy_timeout(Duration::from_secs(5))
}
/// Connect options for one user's database.
///
/// When `key` is present the pool is a SQLCipher pool: sqlx puts `PRAGMA key`
/// ahead of every other pragma, so the page cipher is armed before `journal_mode`
/// touches the file. Without a key the same code opens an ordinary SQLite file —
/// which is exactly what a cleartext user gets.
///
/// The returned value carries the raw key. **Never** `{:?}` it: `SqliteConnectOptions`
/// prints its pragmas, so a debug format anywhere near this would write the DEK,
/// in hex, into the logs.
fn user_options(path: &Path, key: Option<&Dek>, create: bool) -> SqliteConnectOptions {
let opts = SqliteConnectOptions::new().filename(path).create_if_missing(create);
let opts = match key {
Some(dek) => opts.pragma("key", dek.to_pragma()),
None => opts,
};
tuned(opts)
}
/// Forces a page read so a wrong key fails here, at open time, rather than at the
/// first unrelated query. SQLCipher answers "file is not a database" when the key
/// does not decrypt the header.
async fn probe(pool: &SqlitePool) -> Result<()> {
sqlx::query_scalar::<_, i64>("SELECT count(*) FROM sqlite_master")
.fetch_one(pool)
.await
.context("database did not open — wrong key, or not a database")?;
Ok(())
}
/// The system database: registry tables, plus — for now — the owner tables.
///
/// Owner tables live here **transitionally**. Nothing has been migrated to
/// per-user pools yet, so sessions, history and jobs still land in `system.db`.
/// When the call sites move to `UserManager::pool_of`, this second call goes
/// away and the owner-without-a-user gets a file of its own.
pub async fn init_system_pool(path: &str) -> Result<SqlitePool> {
ensure_parent(Path::new(path))?;
let opts = tuned(SqliteConnectOptions::from_str(path)?.create_if_missing(true));
let pool = SqlitePool::connect_with(opts).await?;
create_registry_tables(&pool).await?;
create_owner_tables(&pool).await?;
crate::boot::section("Database initialised".to_string());
Ok(pool)
}
/// Provisions `database/{userid}.db` and lays down the owner schema.
///
/// Only this function may create a user's database. Login goes through
/// [`open_user_pool`], which refuses to create anything: a missing file there is
/// data loss, and creating a fresh empty one under the right password would hide
/// it instead of reporting it.
pub async fn create_user_pool(path: &Path, key: Option<&Dek>) -> Result<SqlitePool> {
ensure_parent(path)?;
let pool = SqlitePool::connect_with(user_options(path, key, true)).await?;
probe(&pool).await?;
create_owner_tables(&pool).await?;
Ok(pool)
}
/// Opens an existing user database. Never creates one — see [`create_user_pool`].
pub async fn open_user_pool(path: &Path, key: Option<&Dek>) -> Result<SqlitePool> {
let pool = SqlitePool::connect_with(user_options(path, key, false)).await?;
probe(&pool).await?;
Ok(pool)
}
// ── Registry tables ───────────────────────────────────────────────────────────
//
// Instance-wide, readable without any user key: the directory you must open
// before you know who exists. Nothing here is scoped to one user.
async fn create_registry_tables(pool: &SqlitePool) -> Result<()> {
sqlx::query(
"CREATE TABLE IF NOT EXISTS llm_providers (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL UNIQUE,
type TEXT NOT NULL,
api_key TEXT,
base_url TEXT,
description TEXT,
removed_at TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
// `name` is the unique identity + resolution key (LlmManager keys its
// in-memory model map by name). There is deliberately NO
// UNIQUE(provider_id, model_id): the same underlying model may be
// registered multiple times under one provider with different aliases
// and reasoning settings (e.g. "glm-4.6" vs "glm-4.6-thinking").
"CREATE TABLE IF NOT EXISTS llm_models (
id INTEGER PRIMARY KEY AUTOINCREMENT,
provider_id INTEGER NOT NULL REFERENCES llm_providers(id) ON DELETE CASCADE,
model_id TEXT NOT NULL,
name TEXT NOT NULL UNIQUE,
strength TEXT,
scope TEXT NOT NULL DEFAULT '[]',
is_default INTEGER NOT NULL DEFAULT 0,
priority INTEGER NOT NULL DEFAULT 100,
extra_params TEXT,
removed_at TEXT,
context_length INTEGER,
max_output_tokens INTEGER,
knowledge_cutoff TEXT,
capabilities TEXT NOT NULL DEFAULT '[]',
reasoning TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS transcribe_models (
id INTEGER PRIMARY KEY AUTOINCREMENT,
provider_id INTEGER NOT NULL REFERENCES llm_providers(id),
model_id TEXT NOT NULL,
name TEXT NOT NULL UNIQUE,
language TEXT,
priority INTEGER NOT NULL DEFAULT 100,
removed_at TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE(provider_id, model_id)
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS image_generate_models (
id INTEGER PRIMARY KEY AUTOINCREMENT,
provider_id INTEGER NOT NULL REFERENCES llm_providers(id),
model_id TEXT NOT NULL,
name TEXT NOT NULL UNIQUE,
priority INTEGER NOT NULL DEFAULT 100,
removed_at TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE(provider_id, model_id)
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS tts_models (
id INTEGER PRIMARY KEY AUTOINCREMENT,
provider_id INTEGER NOT NULL REFERENCES llm_providers(id),
model_id TEXT NOT NULL,
voice_id TEXT,
name TEXT NOT NULL UNIQUE,
description TEXT,
instructions TEXT,
priority INTEGER NOT NULL DEFAULT 100,
removed_at TEXT,
response_format TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE(provider_id, model_id)
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS plugins (
id TEXT PRIMARY KEY,
enabled INTEGER NOT NULL DEFAULT 0,
config TEXT NOT NULL DEFAULT '{}',
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS tool_permission_groups (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
description TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS approval_rules (
id INTEGER PRIMARY KEY AUTOINCREMENT,
agent_id TEXT,
source TEXT,
tool_pattern TEXT NOT NULL,
action TEXT NOT NULL DEFAULT 'require'
CHECK(action IN ('require', 'allow', 'deny')),
note TEXT,
priority INTEGER NOT NULL DEFAULT 100,
path_pattern TEXT,
group_id TEXT REFERENCES tool_permission_groups(id),
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS config (
key TEXT PRIMARY KEY,
value TEXT NOT NULL,
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
// Every tool ever offered to the LLM, recorded by `ToolDiscovery` at
// injection time. Lets the approval / Security-groups UI list and gate tools
// that are injected dynamically outside the `ToolRegistry` (interface tools,
// plugin tools, provider tools).
sqlx::query(
"CREATE TABLE IF NOT EXISTS known_tools (
name TEXT PRIMARY KEY,
description TEXT NOT NULL DEFAULT '',
schema TEXT,
first_seen INTEGER NOT NULL DEFAULT (strftime('%s','now')),
last_seen INTEGER NOT NULL DEFAULT (strftime('%s','now'))
)",
)
.execute(pool)
.await?;
// Spend/telemetry metadata. Stays in the registry so cost charts run without
// decrypting anything: the admin sees how much, when and which model — never
// what was said. `session_id` / `stack_id` are bare integers, not foreign
// keys, precisely because the rows they point at live in another file.
sqlx::query(
"CREATE TABLE IF NOT EXISTS llm_requests (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id INTEGER,
stack_id INTEGER,
model_name TEXT NOT NULL,
request_json TEXT NOT NULL DEFAULT '',
request_headers TEXT,
response_json TEXT,
response_headers TEXT,
error_text TEXT,
input_tokens INTEGER,
output_tokens INTEGER,
cache_read_tokens INTEGER,
cache_creation_tokens INTEGER,
duration_ms INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_llm_requests_created
ON llm_requests (created_at)",
)
.execute(pool)
.await?;
// User directory + auth material. Read before every login, so it lives in
// the registry — which means it must never hold anything that derives a
// user's key: `database_password` is the DEK sealed under a key derived
// from the password, useless without it.
//
// `role_id` has no `REFERENCES roles(id)` yet: sqlx turns on
// `PRAGMA foreign_keys`, so pointing at a table that does not exist would
// make every INSERT fail. The constraint lands with the `roles` table.
sqlx::query(
"CREATE TABLE IF NOT EXISTS users (
id TEXT PRIMARY KEY,
username TEXT NOT NULL UNIQUE,
display_name TEXT,
role_id TEXT NOT NULL,
encrypted INTEGER NOT NULL,
kdf_params TEXT,
kdf_salt BLOB,
database_password BLOB,
password_hash BLOB,
active INTEGER NOT NULL DEFAULT 1,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
CHECK (
(encrypted = 1 AND database_password IS NOT NULL AND password_hash IS NULL)
OR (encrypted = 0 AND database_password IS NULL)
)
)",
)
.execute(pool)
.await?;
Ok(())
}
// ── Owner tables ──────────────────────────────────────────────────────────────
//
// One owner's content. The schema is identical in every file that has it — only
// the file says who owns the rows — so this runs verbatim against `system.db`
// and against each `database/{userid}.db`.
//
// **No foreign key here may point at a registry table.** SQLite cannot enforce a
// key across files, not even through `ATTACH`, and sqlx turns on
// `PRAGMA foreign_keys`: the `CREATE TABLE` would succeed and every `INSERT`
// would fail. `create_owner_tables_stand_alone` in the tests guards this by
// running the schema against a file that has nothing else in it.
pub async fn create_owner_tables(pool: &SqlitePool) -> Result<()> {
sqlx::query(
"CREATE TABLE IF NOT EXISTS chat_sessions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
title TEXT,
source TEXT NOT NULL DEFAULT 'web',
agent_id TEXT NOT NULL DEFAULT 'main',
is_interactive INTEGER NOT NULL DEFAULT 1,
is_ephemeral INTEGER NOT NULL DEFAULT 0,
run_context TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS chat_sessions_stack (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id INTEGER NOT NULL REFERENCES chat_sessions(id),
agent_id TEXT NOT NULL DEFAULT 'main',
agent_prompt TEXT,
depth INTEGER NOT NULL DEFAULT 0,
parent_tool_call_id INTEGER,
terminated_at TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
// `model_db_id` used to point at `llm_models(id)`. It was write-only — never
// selected, never joined — and it was the one key that crossed into the
// registry. Which model answered is already recorded in
// `llm_requests.model_name`.
sqlx::query(
"CREATE TABLE IF NOT EXISTS chat_history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_stack_id INTEGER NOT NULL REFERENCES chat_sessions_stack(id),
role TEXT NOT NULL CHECK(role IN ('user', 'assistant', 'agent')),
content TEXT NOT NULL DEFAULT '',
status TEXT NOT NULL DEFAULT 'ok' CHECK(status IN ('ok', 'failed')),
input_tokens INTEGER,
output_tokens INTEGER,
duration_ms INTEGER,
is_synthetic INTEGER NOT NULL DEFAULT 0,
reasoning_content TEXT,
cost REAL,
metadata TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS chat_llm_tools (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL REFERENCES chat_history(id),
name TEXT NOT NULL,
arguments TEXT,
result TEXT,
status TEXT NOT NULL DEFAULT 'running' CHECK(status IN ('running', 'pending', 'done', 'failed', 'cancelled', 'rejected')),
result_type TEXT NOT NULL DEFAULT 'string' CHECK(result_type IN ('string', 'json')),
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_stack_session ON chat_sessions_stack(session_id)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_history_stack ON chat_history(session_stack_id)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_tools_message ON chat_llm_tools(message_id)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS chat_summaries (
id INTEGER PRIMARY KEY AUTOINCREMENT,
stack_id INTEGER NOT NULL REFERENCES chat_sessions_stack(id),
content TEXT NOT NULL,
covers_up_to_message_id INTEGER NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_chat_summaries_stack
ON chat_summaries (stack_id)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS session_scratchpad (
session_id INTEGER NOT NULL REFERENCES chat_sessions(id),
key TEXT NOT NULL,
value TEXT NOT NULL,
PRIMARY KEY (session_id, key)
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS session_mcp_grants (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id INTEGER NOT NULL,
mcp_name TEXT NOT NULL,
granted_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE(session_id, mcp_name)
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS stack_mcp_grants (
id INTEGER PRIMARY KEY AUTOINCREMENT,
stack_id INTEGER NOT NULL,
mcp_name TEXT NOT NULL,
granted_at TEXT NOT NULL DEFAULT (datetime('now')),
UNIQUE(stack_id, mcp_name)
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS scheduled_jobs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
title TEXT NOT NULL,
description TEXT NOT NULL DEFAULT '',
cron TEXT NOT NULL,
prompt TEXT NOT NULL,
agent_id TEXT NOT NULL DEFAULT 'main',
session_id INTEGER REFERENCES chat_sessions(id),
enabled INTEGER NOT NULL DEFAULT 1,
last_run_at TEXT,
next_run_at TEXT,
single_run INTEGER NOT NULL DEFAULT 0,
running_session_id INTEGER,
kind TEXT NOT NULL DEFAULT 'cron',
parent_session_id INTEGER REFERENCES chat_sessions(id),
run_context TEXT,
running_since TEXT,
origin_ref TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS job_runs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
job_id INTEGER NOT NULL REFERENCES scheduled_jobs(id),
session_id INTEGER,
started_at TEXT NOT NULL,
completed_at TEXT,
duration_ms INTEGER,
status TEXT NOT NULL
CHECK(status IN ('completed', 'failed', 'cancelled')),
final_response TEXT,
error TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_job_runs_job_id
ON job_runs (job_id, created_at DESC)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS mcp_servers (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL UNIQUE,
transport TEXT NOT NULL DEFAULT 'stdio',
command TEXT,
args_json TEXT,
env_json TEXT,
url TEXT,
api_key TEXT,
description TEXT,
friendly_name TEXT,
enabled INTEGER NOT NULL DEFAULT 1,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS mcp_events (
id INTEGER PRIMARY KEY AUTOINCREMENT,
source TEXT NOT NULL,
method TEXT NOT NULL,
payload TEXT NOT NULL,
processed INTEGER NOT NULL DEFAULT 0,
processed_at TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_mcp_events_pending
ON mcp_events (processed, created_at)
WHERE processed = 0",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS sources (
id TEXT PRIMARY KEY,
active_session_id INTEGER REFERENCES chat_sessions(id),
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS secrets (
key TEXT PRIMARY KEY,
value TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS projects (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
path TEXT NOT NULL,
description TEXT NOT NULL DEFAULT '',
run_context TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS project_tickets (
id INTEGER PRIMARY KEY AUTOINCREMENT,
project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
title TEXT NOT NULL,
description TEXT NOT NULL DEFAULT '',
status TEXT NOT NULL DEFAULT 'todo'
CHECK(status IN ('todo','pending','in_progress','done','failed')),
agent_id TEXT NOT NULL DEFAULT 'main',
run_context TEXT,
job_id INTEGER REFERENCES scheduled_jobs(id),
result TEXT,
error TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
started_at TEXT,
completed_at TEXT
)",
)
.execute(pool)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn temp_dir(tag: &str) -> PathBuf {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
let mut p = std::env::temp_dir();
p.push(format!("skald-db-{tag}-{}-{nanos}", std::process::id()));
p
}
/// The guardrail the file split would have given for free.
///
/// `create_owner_tables` runs here against a database that holds *nothing
/// else* — no `llm_models`, no `users`. Every table then takes a row. Any
/// foreign key that reaches into a registry table passes `CREATE TABLE` and
/// dies on the `INSERT`, right here, instead of lying dormant until a real
/// user logs in and writes to their own file.
#[tokio::test]
async fn owner_tables_stand_alone_with_foreign_keys_on() {
let dir = temp_dir("owner-standalone");
let path = dir.join("owner.db");
let pool = create_user_pool(&path, None).await.unwrap();
let (fk,): (i64,) = sqlx::query_as("PRAGMA foreign_keys")
.fetch_one(&pool).await.unwrap();
assert_eq!(fk, 1, "the guardrail is meaningless without FK enforcement");
let one = |q: &'static str| sqlx::query(q).execute(&pool);
one("INSERT INTO chat_sessions (id, title) VALUES (1, 't')").await.unwrap();
one("INSERT INTO chat_sessions_stack (id, session_id) VALUES (1, 1)").await.unwrap();
one("INSERT INTO chat_history (id, session_stack_id, role, content) VALUES (1, 1, 'user', 'hi')")
.await.unwrap();
one("INSERT INTO chat_llm_tools (message_id, name) VALUES (1, 'exec')").await.unwrap();
one("INSERT INTO chat_summaries (stack_id, content, covers_up_to_message_id) VALUES (1, 's', 1)")
.await.unwrap();
one("INSERT INTO session_scratchpad (session_id, key, value) VALUES (1, 'k', 'v')").await.unwrap();
one("INSERT INTO session_mcp_grants (session_id, mcp_name) VALUES (1, 'm')").await.unwrap();
one("INSERT INTO stack_mcp_grants (stack_id, mcp_name) VALUES (1, 'm')").await.unwrap();
one("INSERT INTO scheduled_jobs (id, title, cron, prompt, session_id) VALUES (1, 't', '* * * * *', 'p', 1)")
.await.unwrap();
one("INSERT INTO job_runs (job_id, started_at, status) VALUES (1, 'now', 'completed')").await.unwrap();
one("INSERT INTO mcp_servers (name) VALUES ('srv')").await.unwrap();
one("INSERT INTO mcp_events (source, method, payload) VALUES ('s', 'm', '{}')").await.unwrap();
one("INSERT INTO sources (id, active_session_id) VALUES ('web', 1)").await.unwrap();
one("INSERT INTO secrets (key, value) VALUES ('k', 'v')").await.unwrap();
one("INSERT INTO projects (id, name, path) VALUES (1, 'p', '/tmp')").await.unwrap();
one("INSERT INTO project_tickets (project_id, title, job_id) VALUES (1, 't', 1)").await.unwrap();
pool.close().await;
let _ = std::fs::remove_dir_all(&dir);
}
/// A cleartext user's database is an ordinary SQLite file; an encrypted one
/// is unreadable without the key, and does not even carry SQLite's header.
#[tokio::test]
async fn an_encrypted_database_is_opaque_without_its_key() {
let dir = temp_dir("opaque");
let clear_path = dir.join("clear.db");
let enc_path = dir.join("enc.db");
let clear = create_user_pool(&clear_path, None).await.unwrap();
clear.close().await;
let dek = Dek::random();
let enc = create_user_pool(&enc_path, Some(&dek)).await.unwrap();
sqlx::query("INSERT INTO chat_sessions (title) VALUES ('secret')")
.execute(&enc).await.unwrap();
enc.close().await;
let magic = b"SQLite format 3\0";
assert_eq!(&std::fs::read(&clear_path).unwrap()[..16], magic);
assert_ne!(&std::fs::read(&enc_path).unwrap()[..16], magic,
"an encrypted file must not advertise itself as SQLite");
assert!(open_user_pool(&enc_path, None).await.is_err(), "no key must not open it");
assert!(open_user_pool(&enc_path, Some(&Dek::random())).await.is_err(), "wrong key must not open it");
// ...and the right key still does, with the row intact.
let reopened = open_user_pool(&enc_path, Some(&dek)).await.unwrap();
let (title,): (String,) = sqlx::query_as("SELECT title FROM chat_sessions")
.fetch_one(&reopened).await.unwrap();
assert_eq!(title, "secret");
reopened.close().await;
let _ = std::fs::remove_dir_all(&dir);
}
/// Login must never conjure a database. A missing file is data loss and has
/// to be reported, not papered over with a fresh empty one.
#[tokio::test]
async fn open_never_creates_a_database() {
let dir = temp_dir("no-create");
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("ghost.db");
assert!(open_user_pool(&path, None).await.is_err());
assert!(open_user_pool(&path, Some(&Dek::random())).await.is_err());
assert!(!path.exists(), "open_user_pool must not have created the file");
let _ = std::fs::remove_dir_all(&dir);
}
}
+49
View File
@@ -0,0 +1,49 @@
use anyhow::Result;
use sqlx::SqlitePool;
#[derive(Debug, Clone)]
pub struct PluginRow {
pub id: String,
pub enabled: bool,
pub config: String, // JSON blob
}
/// Returns (enabled, config_json) for a plugin, or None if not yet in DB.
pub async fn get(pool: &SqlitePool, id: &str) -> Result<Option<PluginRow>> {
let row: Option<(String, i64, String)> = sqlx::query_as(
"SELECT id, enabled, config FROM plugins WHERE id = ?1",
)
.bind(id)
.fetch_optional(pool)
.await?;
Ok(row.map(|(id, e, config)| PluginRow { id, enabled: e != 0, config }))
}
/// Upserts both enabled flag and config JSON.
pub async fn upsert(pool: &SqlitePool, id: &str, enabled: bool, config: &str) -> Result<()> {
sqlx::query(
"INSERT INTO plugins (id, enabled, config)
VALUES (?1, ?2, ?3)
ON CONFLICT(id) DO UPDATE SET enabled = excluded.enabled,
config = excluded.config",
)
.bind(id)
.bind(enabled as i64)
.bind(config)
.execute(pool)
.await?;
Ok(())
}
/// Returns all plugin rows. Used by the config watcher.
pub async fn list(pool: &SqlitePool) -> Result<Vec<PluginRow>> {
let rows: Vec<(String, i64, String)> = sqlx::query_as(
"SELECT id, enabled, config FROM plugins ORDER BY id",
)
.fetch_all(pool)
.await?;
Ok(rows
.into_iter()
.map(|(id, e, config)| PluginRow { id, enabled: e != 0, config })
.collect())
}
+149
View File
@@ -0,0 +1,149 @@
use anyhow::Result;
use sqlx::SqlitePool;
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct ProjectTicket {
pub id: i64,
pub project_id: i64,
pub title: String,
pub description: String,
pub status: String,
pub agent_id: String,
pub run_context: Option<String>,
pub job_id: Option<i64>,
pub result: Option<String>,
pub error: Option<String>,
pub created_at: String,
pub started_at: Option<String>,
pub completed_at: Option<String>,
pub session_id: Option<i64>,
}
const SELECT: &str =
"SELECT pt.id, pt.project_id, pt.title, pt.description, pt.status, pt.agent_id,
pt.run_context, pt.job_id, pt.result, pt.error, pt.created_at,
pt.started_at, pt.completed_at,
COALESCE(sj.running_session_id,
(SELECT session_id FROM job_runs
WHERE job_id = pt.job_id ORDER BY id DESC LIMIT 1)
) AS session_id
FROM project_tickets pt
LEFT JOIN scheduled_jobs sj ON sj.id = pt.job_id";
pub async fn list_for_project(pool: &SqlitePool, project_id: i64) -> Result<Vec<ProjectTicket>> {
let rows = sqlx::query_as::<_, ProjectTicket>(sqlx::AssertSqlSafe(format!(
"{SELECT} WHERE pt.project_id = ? ORDER BY pt.id"
)))
.bind(project_id)
.fetch_all(pool)
.await?;
Ok(rows)
}
pub async fn get(pool: &SqlitePool, id: i64) -> Result<Option<ProjectTicket>> {
let row = sqlx::query_as::<_, ProjectTicket>(sqlx::AssertSqlSafe(format!(
"{SELECT} WHERE pt.id = ?"
)))
.bind(id)
.fetch_optional(pool)
.await?;
Ok(row)
}
pub async fn create(
pool: &SqlitePool,
project_id: i64,
title: &str,
description: &str,
agent_id: &str,
run_context: Option<&str>,
) -> Result<ProjectTicket> {
let id = sqlx::query(
"INSERT INTO project_tickets (project_id, title, description, agent_id, run_context)
VALUES (?, ?, ?, ?, ?)",
)
.bind(project_id)
.bind(title)
.bind(description)
.bind(agent_id)
.bind(run_context)
.execute(pool)
.await?
.last_insert_rowid();
let row = sqlx::query_as::<_, ProjectTicket>(sqlx::AssertSqlSafe(format!(
"{SELECT} WHERE pt.id = ?"
)))
.bind(id)
.fetch_one(pool)
.await?;
Ok(row)
}
pub async fn delete(pool: &SqlitePool, id: i64) -> Result<bool> {
let n = sqlx::query("DELETE FROM project_tickets WHERE id = ?")
.bind(id)
.execute(pool)
.await?
.rows_affected();
Ok(n > 0)
}
pub async fn set_status(pool: &SqlitePool, id: i64, status: &str) -> Result<()> {
sqlx::query("UPDATE project_tickets SET status = ? WHERE id = ?")
.bind(status)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
/// Mark as in_progress and record the scheduled job that is running it.
pub async fn start(pool: &SqlitePool, id: i64, job_id: i64) -> Result<()> {
sqlx::query(
"UPDATE project_tickets
SET status = 'in_progress', job_id = ?, started_at = datetime('now')
WHERE id = ?",
)
.bind(job_id)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
/// Mark as done or failed, recording result/error and timestamp.
pub async fn complete(
pool: &SqlitePool,
id: i64,
result: Option<&str>,
error: Option<&str>,
) -> Result<()> {
let status = if error.is_some() { "failed" } else { "done" };
sqlx::query(
"UPDATE project_tickets
SET status = ?, result = ?, error = ?, completed_at = datetime('now')
WHERE id = ?",
)
.bind(status)
.bind(result)
.bind(error)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
/// Reset a ticket back to todo, clearing all run state.
pub async fn reset(pool: &SqlitePool, id: i64) -> Result<()> {
sqlx::query(
"UPDATE project_tickets
SET status = 'todo', job_id = NULL, result = NULL, error = NULL,
started_at = NULL, completed_at = NULL
WHERE id = ?",
)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
+103
View File
@@ -0,0 +1,103 @@
use anyhow::Result;
use sqlx::SqlitePool;
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct Project {
pub id: i64,
pub name: String,
pub path: String,
pub description: String,
pub run_context: Option<String>,
pub created_at: String,
pub updated_at: String,
}
const SELECT: &str =
"SELECT id, name, path, description, run_context, created_at, updated_at
FROM projects";
pub async fn list(pool: &SqlitePool) -> Result<Vec<Project>> {
let rows = sqlx::query_as::<_, Project>(sqlx::AssertSqlSafe(format!(
"{SELECT} ORDER BY updated_at DESC"
)))
.fetch_all(pool)
.await?;
Ok(rows)
}
pub async fn get(pool: &SqlitePool, id: i64) -> Result<Option<Project>> {
let row = sqlx::query_as::<_, Project>(sqlx::AssertSqlSafe(format!("{SELECT} WHERE id = ?")))
.bind(id)
.fetch_optional(pool)
.await?;
Ok(row)
}
pub async fn create(
pool: &SqlitePool,
name: &str,
path: &str,
description: &str,
run_context: Option<&str>,
) -> Result<Project> {
let id = sqlx::query(
"INSERT INTO projects (name, path, description, run_context)
VALUES (?, ?, ?, ?)",
)
.bind(name)
.bind(path)
.bind(description)
.bind(run_context)
.execute(pool)
.await?
.last_insert_rowid();
let row = sqlx::query_as::<_, Project>(sqlx::AssertSqlSafe(format!("{SELECT} WHERE id = ?")))
.bind(id)
.fetch_one(pool)
.await?;
Ok(row)
}
pub async fn update(
pool: &SqlitePool,
id: i64,
name: &str,
path: &str,
description: &str,
run_context: Option<&str>,
) -> Result<bool> {
let n = sqlx::query(
"UPDATE projects
SET name = ?, path = ?, description = ?, run_context = ?,
updated_at = datetime('now')
WHERE id = ?",
)
.bind(name)
.bind(path)
.bind(description)
.bind(run_context)
.bind(id)
.execute(pool)
.await?
.rows_affected();
Ok(n > 0)
}
/// Touch updated_at — called after every ticket operation so ordering by recency works.
pub async fn touch(pool: &SqlitePool, id: i64) -> Result<()> {
sqlx::query("UPDATE projects SET updated_at = datetime('now') WHERE id = ?")
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn delete(pool: &SqlitePool, id: i64) -> Result<bool> {
let n = sqlx::query("DELETE FROM projects WHERE id = ?")
.bind(id)
.execute(pool)
.await?
.rows_affected();
Ok(n > 0)
}
+211
View File
@@ -0,0 +1,211 @@
use anyhow::Result;
use sqlx::SqlitePool;
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct ScheduledJob {
pub id: i64,
pub title: String,
pub description: String,
pub cron: String,
pub prompt: String,
pub agent_id: String,
pub session_id: Option<i64>,
pub enabled: bool,
pub last_run_at: Option<String>,
pub next_run_at: Option<String>,
pub single_run: bool,
pub running_session_id: Option<i64>,
pub running_since: Option<String>,
pub kind: String,
pub created_at: String,
pub parent_session_id: Option<i64>,
pub run_context: Option<String>,
pub origin_ref: Option<String>,
}
const SELECT: &str =
"SELECT id, title, description, cron, prompt, agent_id, session_id,
CAST(enabled AS BOOLEAN) AS enabled,
last_run_at,
next_run_at,
CAST(single_run AS BOOLEAN) AS single_run,
running_session_id,
running_since,
kind,
created_at,
parent_session_id,
run_context,
origin_ref
FROM scheduled_jobs";
pub async fn get_by_id(pool: &SqlitePool, id: i64) -> Result<Option<ScheduledJob>> {
sqlx::query_as::<_, ScheduledJob>(sqlx::AssertSqlSafe(format!("{SELECT} WHERE id = ?")))
.bind(id)
.fetch_optional(pool)
.await
.map_err(Into::into)
}
pub async fn list(pool: &SqlitePool) -> Result<Vec<ScheduledJob>> {
let rows = sqlx::query_as::<_, ScheduledJob>(sqlx::AssertSqlSafe(format!("{SELECT} ORDER BY id")))
.fetch_all(pool)
.await?;
Ok(rows)
}
/// Jobs enabled and due to run: next_run_at is in the past and not currently running.
/// `now_rfc3339` should be `chrono::Utc::now().to_rfc3339()`.
pub async fn list_due(pool: &SqlitePool, now_rfc3339: &str) -> Result<Vec<ScheduledJob>> {
let rows = sqlx::query_as::<_, ScheduledJob>(sqlx::AssertSqlSafe(format!(
"{SELECT}
WHERE kind = 'cron'
AND enabled = 1
AND next_run_at IS NOT NULL
AND next_run_at <= ?
AND running_session_id IS NULL
ORDER BY next_run_at",
)))
.bind(now_rfc3339)
.fetch_all(pool)
.await?;
Ok(rows)
}
/// Jobs that were running when the process was last killed (running_session_id IS NOT NULL).
pub async fn list_interrupted(pool: &SqlitePool) -> Result<Vec<ScheduledJob>> {
let rows = sqlx::query_as::<_, ScheduledJob>(sqlx::AssertSqlSafe(format!(
"{SELECT} WHERE running_session_id IS NOT NULL ORDER BY id",
)))
.fetch_all(pool)
.await?;
Ok(rows)
}
pub async fn create(
pool: &SqlitePool,
title: &str,
description: &str,
cron: &str,
prompt: &str,
agent_id: &str,
single_run: bool,
next_run_at: Option<&str>,
kind: &str,
parent_session_id: Option<i64>,
run_context: Option<&str>,
origin_ref: Option<&str>,
) -> Result<ScheduledJob> {
let id = sqlx::query(
"INSERT INTO scheduled_jobs (title, description, cron, prompt, agent_id, single_run, next_run_at, kind, parent_session_id, run_context, origin_ref)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind(title)
.bind(description)
.bind(cron)
.bind(prompt)
.bind(agent_id)
.bind(single_run as i64)
.bind(next_run_at)
.bind(kind)
.bind(parent_session_id)
.bind(run_context)
.bind(origin_ref)
.execute(pool)
.await?
.last_insert_rowid();
let row = sqlx::query_as::<_, ScheduledJob>(sqlx::AssertSqlSafe(format!("{SELECT} WHERE id = ?")))
.bind(id)
.fetch_one(pool)
.await?;
Ok(row)
}
pub async fn delete(pool: &SqlitePool, id: i64) -> Result<bool> {
// Clear the soft back-reference from project_tickets first: its job_id FK has
// no ON DELETE action, so a ticket still pointing at this job would block the
// scheduled_jobs DELETE with a FOREIGN KEY constraint failure.
sqlx::query("UPDATE project_tickets SET job_id = NULL WHERE job_id = ?")
.bind(id)
.execute(pool)
.await?;
sqlx::query("DELETE FROM job_runs WHERE job_id = ?")
.bind(id)
.execute(pool)
.await?;
let n = sqlx::query("DELETE FROM scheduled_jobs WHERE id = ?")
.bind(id)
.execute(pool)
.await?
.rows_affected();
Ok(n > 0)
}
pub async fn set_enabled(pool: &SqlitePool, id: i64, enabled: bool) -> Result<bool> {
let n = sqlx::query("UPDATE scheduled_jobs SET enabled = ? WHERE id = ?")
.bind(enabled as i64)
.bind(id)
.execute(pool)
.await?
.rows_affected();
Ok(n > 0)
}
/// Update next_run_at without touching anything else (used when re-enabling a job).
pub async fn set_next_run_at(pool: &SqlitePool, id: i64, next_run_at: &str) -> Result<()> {
sqlx::query("UPDATE scheduled_jobs SET next_run_at = ? WHERE id = ?")
.bind(next_run_at)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
/// Mark a job as in-flight. Called at the start of run_job(), before handle_message().
pub async fn set_running(pool: &SqlitePool, id: i64, session_id: i64) -> Result<()> {
sqlx::query(
"UPDATE scheduled_jobs SET running_session_id = ?, running_since = datetime('now') WHERE id = ?",
)
.bind(session_id)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
pub async fn set_run_context(pool: &SqlitePool, id: i64, run_context: Option<&str>) -> Result<bool> {
let n = sqlx::query("UPDATE scheduled_jobs SET run_context = ? WHERE id = ?")
.bind(run_context)
.bind(id)
.execute(pool)
.await?
.rows_affected();
Ok(n > 0)
}
/// Mark a job as finished. Called at the end of run_job() regardless of outcome.
///
/// - Sets `last_run_at = now`, clears `running_session_id`.
/// - If `next_run_at` is `Some`: updates the field (next scheduled fire).
/// - If `next_run_at` is `None` (single-run job): sets `enabled = 0`.
pub async fn finish_run(
pool: &SqlitePool,
id: i64,
next_run_at: Option<&str>,
) -> Result<()> {
sqlx::query(
"UPDATE scheduled_jobs
SET last_run_at = datetime('now'),
running_session_id = NULL,
running_since = NULL,
next_run_at = COALESCE(?, next_run_at),
enabled = CASE WHEN ? IS NULL THEN 0 ELSE enabled END
WHERE id = ?",
)
.bind(next_run_at)
.bind(next_run_at)
.bind(id)
.execute(pool)
.await?;
Ok(())
}
+26
View File
@@ -0,0 +1,26 @@
use anyhow::Result;
use sqlx::SqlitePool;
pub async fn upsert(pool: &SqlitePool, session_id: i64, key: &str, value: &str) -> Result<()> {
sqlx::query(
"INSERT INTO session_scratchpad (session_id, key, value)
VALUES (?, ?, ?)
ON CONFLICT (session_id, key) DO UPDATE SET value = excluded.value"
)
.bind(session_id)
.bind(key)
.bind(value)
.execute(pool)
.await?;
Ok(())
}
pub async fn for_session(pool: &SqlitePool, session_id: i64) -> Result<Vec<(String, String)>> {
let rows = sqlx::query_as::<_, (String, String)>(
"SELECT key, value FROM session_scratchpad WHERE session_id = ? ORDER BY key"
)
.bind(session_id)
.fetch_all(pool)
.await?;
Ok(rows)
}
@@ -0,0 +1,36 @@
use anyhow::Result;
use sqlx::SqlitePool;
/// Grant access to an MCP server for a session.
/// Uses INSERT OR IGNORE so calling it multiple times is safe.
pub async fn grant(pool: &SqlitePool, session_id: i64, mcp_name: &str) -> Result<()> {
sqlx::query(
"INSERT OR IGNORE INTO session_mcp_grants (session_id, mcp_name)
VALUES (?, ?)"
)
.bind(session_id)
.bind(mcp_name)
.execute(pool)
.await?;
Ok(())
}
/// Revoke all MCP grants for a session.
pub async fn revoke_all(pool: &SqlitePool, session_id: i64) -> Result<()> {
sqlx::query("DELETE FROM session_mcp_grants WHERE session_id = ?")
.bind(session_id)
.execute(pool)
.await?;
Ok(())
}
/// Returns the names of all MCP servers granted for this session.
pub async fn list_for_session(pool: &SqlitePool, session_id: i64) -> Result<Vec<String>> {
let rows = sqlx::query_as::<_, (String,)>(
"SELECT mcp_name FROM session_mcp_grants WHERE session_id = ? ORDER BY granted_at"
)
.bind(session_id)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(|(name,)| name).collect())
}
+47
View File
@@ -0,0 +1,47 @@
use sqlx::SqlitePool;
pub struct Source {
pub id: String,
pub active_session_id: Option<i64>,
pub updated_at: String,
}
/// Upsert a source, setting its active session.
pub async fn upsert(pool: &SqlitePool, id: &str, session_id: i64) -> anyhow::Result<()> {
sqlx::query(
"INSERT INTO sources (id, active_session_id, updated_at)
VALUES (?, ?, datetime('now'))
ON CONFLICT(id) DO UPDATE SET
active_session_id = excluded.active_session_id,
updated_at = excluded.updated_at",
)
.bind(id)
.bind(session_id)
.execute(pool)
.await?;
Ok(())
}
/// Find a source by id.
pub async fn find(pool: &SqlitePool, id: &str) -> anyhow::Result<Option<Source>> {
let row = sqlx::query_as::<_, (String, Option<i64>, String)>(
"SELECT id, active_session_id, updated_at FROM sources WHERE id = ?",
)
.bind(id)
.fetch_optional(pool)
.await?;
Ok(row.map(|(id, active_session_id, updated_at)| Source { id, active_session_id, updated_at }))
}
/// Returns the active session id for a source, if set.
pub async fn active_session_id(pool: &SqlitePool, id: &str) -> anyhow::Result<Option<i64>> {
let row = sqlx::query_as::<_, (Option<i64>,)>(
"SELECT active_session_id FROM sources WHERE id = ?",
)
.bind(id)
.fetch_optional(pool)
.await?;
Ok(row.and_then(|(sid,)| sid))
}
@@ -0,0 +1,36 @@
use anyhow::Result;
use sqlx::SqlitePool;
/// Persist an MCP grant scoped to a specific stack frame (sub-agent).
/// Uses INSERT OR IGNORE so calling it multiple times is safe.
pub async fn grant(pool: &SqlitePool, stack_id: i64, mcp_name: &str) -> Result<()> {
sqlx::query(
"INSERT OR IGNORE INTO stack_mcp_grants (stack_id, mcp_name)
VALUES (?, ?)",
)
.bind(stack_id)
.bind(mcp_name)
.execute(pool)
.await?;
Ok(())
}
/// Returns the names of all MCP servers granted for this stack frame.
pub async fn list_for_stack(pool: &SqlitePool, stack_id: i64) -> Result<Vec<String>> {
let rows = sqlx::query_as::<_, (String,)>(
"SELECT mcp_name FROM stack_mcp_grants WHERE stack_id = ? ORDER BY granted_at",
)
.bind(stack_id)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(|(name,)| name).collect())
}
/// Removes all MCP grants for a stack frame. Called when the frame terminates.
pub async fn delete_for_stack(pool: &SqlitePool, stack_id: i64) -> Result<()> {
sqlx::query("DELETE FROM stack_mcp_grants WHERE stack_id = ?")
.bind(stack_id)
.execute(pool)
.await?;
Ok(())
}
@@ -0,0 +1,85 @@
use anyhow::Result;
use serde::{Deserialize, Serialize};
use sqlx::SqlitePool;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolPermissionGroup {
pub id: String,
pub name: String,
pub description: Option<String>,
pub created_at: String,
}
type RawRow = (String, String, Option<String>, String);
fn from_raw((id, name, description, created_at): RawRow) -> ToolPermissionGroup {
ToolPermissionGroup { id, name, description, created_at }
}
pub async fn list(pool: &SqlitePool) -> Result<Vec<ToolPermissionGroup>> {
let rows = sqlx::query_as::<_, RawRow>(
"SELECT id, name, description, created_at
FROM tool_permission_groups
ORDER BY created_at ASC",
)
.fetch_all(pool)
.await?;
Ok(rows.into_iter().map(from_raw).collect())
}
pub async fn get(pool: &SqlitePool, id: &str) -> Result<Option<ToolPermissionGroup>> {
let row = sqlx::query_as::<_, RawRow>(
"SELECT id, name, description, created_at
FROM tool_permission_groups WHERE id = ?",
)
.bind(id)
.fetch_optional(pool)
.await?;
Ok(row.map(from_raw))
}
pub async fn insert(pool: &SqlitePool, id: &str, name: &str, description: Option<&str>) -> Result<()> {
sqlx::query(
"INSERT INTO tool_permission_groups (id, name, description) VALUES (?, ?, ?)",
)
.bind(id)
.bind(name)
.bind(description)
.execute(pool)
.await?;
Ok(())
}
pub async fn insert_or_ignore(pool: &SqlitePool, id: &str, name: &str, description: Option<&str>) -> Result<()> {
sqlx::query(
"INSERT OR IGNORE INTO tool_permission_groups (id, name, description) VALUES (?, ?, ?)",
)
.bind(id)
.bind(name)
.bind(description)
.execute(pool)
.await?;
Ok(())
}
pub async fn update(pool: &SqlitePool, id: &str, name: &str, description: Option<&str>) -> Result<bool> {
let rows = sqlx::query(
"UPDATE tool_permission_groups SET name = ?, description = ? WHERE id = ?",
)
.bind(name)
.bind(description)
.bind(id)
.execute(pool)
.await?
.rows_affected();
Ok(rows > 0)
}
pub async fn delete(pool: &SqlitePool, id: &str) -> Result<bool> {
let rows = sqlx::query("DELETE FROM tool_permission_groups WHERE id = ?")
.bind(id)
.execute(pool)
.await?
.rows_affected();
Ok(rows > 0)
}
+540
View File
@@ -0,0 +1,540 @@
//! `users` — user directory and auth material.
//!
//! This table lives in the system DB, which anyone owning the box can read. So
//! it must never store anything from which a user's key can be derived.
//!
//! For an **encrypted** user it holds the DEK *wrapped* under a key derived from
//! the password: useless without the password, and the wrap's AEAD tag doubles
//! as the password verifier. That is why an encrypted user has no
//! `password_hash` — a second hash of the same password would only hand an
//! offline attacker an easier target than the wrap itself.
//!
//! A **cleartext** user has no DB key to bind a verifier to, so it carries an
//! ordinary Argon2id hash instead (harmless: that DB is readable anyway).
//!
//! [`Credentials`] makes the two shapes mutually exclusive in the type system,
//! mirroring the `CHECK` constraint on the table.
use anyhow::{Result, anyhow, bail};
use serde::Serialize;
use sqlx::SqlitePool;
/// KDF settings as JSON, e.g. `{"algo":"argon2id","m":65536,"t":3,"p":1}`.
/// Not secret — calibrated on the box when the user is created.
pub type KdfParams = String;
/// Argon2id verifier for a user whose database is not encrypted.
#[derive(Clone)]
pub struct ClearVerifier {
pub kdf_params: KdfParams,
pub kdf_salt: Vec<u8>,
pub password_hash: Vec<u8>,
}
/// Auth material for a user. The variants mirror the table's `CHECK`: an
/// encrypted user has a wrapped DEK and no hash; a cleartext user has no
/// wrapped DEK, and may have no verifier at all (a role that cannot log in).
#[derive(Clone)]
pub enum Credentials {
Encrypted {
kdf_params: KdfParams,
kdf_salt: Vec<u8>,
/// DEK sealed with an AEAD under `KDF(password, kdf_salt)`. Changing the
/// password re-wraps this value; the database itself is never re-encrypted.
database_password: Vec<u8>,
},
Cleartext(Option<ClearVerifier>),
}
impl Credentials {
pub fn is_encrypted(&self) -> bool {
matches!(self, Credentials::Encrypted { .. })
}
}
/// A row of `users`.
///
/// Deliberately **not** `Serialize`: it carries the wrapped DEK and the password
/// verifier, and this type must never be handed to an HTTP handler by accident.
/// Use [`User::summary`] for anything that leaves the process.
#[derive(Clone)]
pub struct User {
pub id: String,
pub username: String,
pub display_name: Option<String>,
pub role_id: String,
pub credentials: Credentials,
pub active: bool,
pub created_at: String,
pub updated_at: String,
}
/// The public-safe projection of a [`User`] — no key material.
#[derive(Debug, Clone, Serialize)]
pub struct UserSummary {
pub id: String,
pub username: String,
pub display_name: Option<String>,
pub role_id: String,
pub encrypted: bool,
pub active: bool,
pub created_at: String,
pub updated_at: String,
}
impl User {
pub fn is_encrypted(&self) -> bool {
self.credentials.is_encrypted()
}
pub fn summary(&self) -> UserSummary {
UserSummary {
id: self.id.clone(),
username: self.username.clone(),
display_name: self.display_name.clone(),
role_id: self.role_id.clone(),
encrypted: self.is_encrypted(),
active: self.active,
created_at: self.created_at.clone(),
updated_at: self.updated_at.clone(),
}
}
}
// Hand-written so a stray `{:?}` — in a tracing span, an error context, a panic
// message — cannot print key material.
impl std::fmt::Debug for Credentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Credentials::Encrypted { .. } => f.write_str("Encrypted(<redacted>)"),
Credentials::Cleartext(None) => f.write_str("Cleartext(no verifier)"),
Credentials::Cleartext(Some(_)) => f.write_str("Cleartext(<redacted>)"),
}
}
}
impl std::fmt::Debug for User {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("User")
.field("id", &self.id)
.field("username", &self.username)
.field("display_name", &self.display_name)
.field("role_id", &self.role_id)
.field("credentials", &self.credentials)
.field("active", &self.active)
.finish()
}
}
// ── Row mapping ───────────────────────────────────────────────────────────────
#[derive(sqlx::FromRow)]
struct Row {
id: String,
username: String,
display_name: Option<String>,
role_id: String,
encrypted: bool,
kdf_params: Option<String>,
kdf_salt: Option<Vec<u8>>,
database_password: Option<Vec<u8>>,
password_hash: Option<Vec<u8>>,
active: bool,
created_at: String,
updated_at: String,
}
/// Builds a `&'static str` (sqlx rejects runtime-built SQL) while keeping the
/// column list — which `Row`'s `FromRow` mirrors — in exactly one place.
macro_rules! select {
($tail:literal) => {
concat!(
"SELECT id, username, display_name, role_id, encrypted, kdf_params, kdf_salt, ",
"database_password, password_hash, active, created_at, updated_at FROM users ",
$tail
)
};
}
impl TryFrom<Row> for User {
type Error = anyhow::Error;
fn try_from(r: Row) -> Result<Self> {
let broken = |what: &str| anyhow!("users row {}: {what}", r.id);
let credentials = if r.encrypted {
Credentials::Encrypted {
kdf_params: r.kdf_params.ok_or_else(|| broken("encrypted without kdf_params"))?,
kdf_salt: r.kdf_salt.ok_or_else(|| broken("encrypted without kdf_salt"))?,
database_password: r.database_password
.ok_or_else(|| broken("encrypted without database_password"))?,
}
} else {
match r.password_hash {
None => Credentials::Cleartext(None),
Some(password_hash) => Credentials::Cleartext(Some(ClearVerifier {
kdf_params: r.kdf_params.ok_or_else(|| broken("verifier without kdf_params"))?,
kdf_salt: r.kdf_salt.ok_or_else(|| broken("verifier without kdf_salt"))?,
password_hash,
})),
}
};
Ok(User {
id: r.id,
username: r.username,
display_name: r.display_name,
role_id: r.role_id,
credentials,
active: r.active,
created_at: r.created_at,
updated_at: r.updated_at,
})
}
}
/// The four credential columns, in table order.
type CredColumns<'a> = (bool, Option<&'a str>, Option<&'a [u8]>, Option<&'a [u8]>, Option<&'a [u8]>);
fn columns(c: &Credentials) -> CredColumns<'_> {
match c {
Credentials::Encrypted { kdf_params, kdf_salt, database_password } => (
true,
Some(kdf_params.as_str()),
Some(kdf_salt.as_slice()),
Some(database_password.as_slice()),
None,
),
Credentials::Cleartext(None) => (false, None, None, None, None),
Credentials::Cleartext(Some(v)) => (
false,
Some(v.kdf_params.as_str()),
Some(v.kdf_salt.as_slice()),
None,
Some(v.password_hash.as_slice()),
),
}
}
// ── Reads ─────────────────────────────────────────────────────────────────────
pub async fn get(pool: &SqlitePool, id: &str) -> Result<Option<User>> {
let row = sqlx::query_as::<_, Row>(select!("WHERE id = ?1"))
.bind(id)
.fetch_optional(pool)
.await?;
row.map(User::try_from).transpose()
}
/// Login entry point: `username` is the handle, `id` is opaque and stable.
pub async fn by_username(pool: &SqlitePool, username: &str) -> Result<Option<User>> {
let row = sqlx::query_as::<_, Row>(select!("WHERE username = ?1"))
.bind(username)
.fetch_optional(pool)
.await?;
row.map(User::try_from).transpose()
}
pub async fn list(pool: &SqlitePool) -> Result<Vec<User>> {
let rows = sqlx::query_as::<_, Row>(select!("ORDER BY username"))
.fetch_all(pool)
.await?;
rows.into_iter().map(User::try_from).collect()
}
pub async fn count(pool: &SqlitePool) -> Result<i64> {
let (n,) = sqlx::query_as::<_, (i64,)>("SELECT COUNT(*) FROM users")
.fetch_one(pool)
.await?;
Ok(n)
}
// ── Writes ────────────────────────────────────────────────────────────────────
/// `id` is supplied by the caller and must be opaque (never the username), so a
/// rename never has to touch `database/{id}.db`.
pub async fn insert(
pool: &SqlitePool,
id: &str,
username: &str,
display_name: Option<&str>,
role_id: &str,
credentials: &Credentials,
) -> Result<()> {
let (encrypted, kdf_params, kdf_salt, database_password, password_hash) = columns(credentials);
sqlx::query(
"INSERT INTO users
(id, username, display_name, role_id, encrypted,
kdf_params, kdf_salt, database_password, password_hash)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
)
.bind(id)
.bind(username)
.bind(display_name)
.bind(role_id)
.bind(encrypted)
.bind(kdf_params)
.bind(kdf_salt)
.bind(database_password)
.bind(password_hash)
.execute(pool)
.await?;
Ok(())
}
/// Replaces the auth material in one statement.
///
/// This is both "change password" (re-wrap the same DEK under a key derived from
/// the new password — the database is never re-encrypted) and the encrypted ↔
/// cleartext migration, since the variant carries the new shape.
pub async fn set_credentials(pool: &SqlitePool, id: &str, credentials: &Credentials) -> Result<()> {
let (encrypted, kdf_params, kdf_salt, database_password, password_hash) = columns(credentials);
let n = sqlx::query(
"UPDATE users SET
encrypted = ?2,
kdf_params = ?3,
kdf_salt = ?4,
database_password = ?5,
password_hash = ?6,
updated_at = datetime('now')
WHERE id = ?1",
)
.bind(id)
.bind(encrypted)
.bind(kdf_params)
.bind(kdf_salt)
.bind(database_password)
.bind(password_hash)
.execute(pool)
.await?
.rows_affected();
if n == 0 {
bail!("no such user: {id}");
}
Ok(())
}
pub async fn set_active(pool: &SqlitePool, id: &str, active: bool) -> Result<()> {
let n = sqlx::query(
"UPDATE users SET active = ?2, updated_at = datetime('now') WHERE id = ?1",
)
.bind(id)
.bind(active)
.execute(pool)
.await?
.rows_affected();
if n == 0 {
bail!("no such user: {id}");
}
Ok(())
}
pub async fn rename(pool: &SqlitePool, id: &str, username: &str, display_name: Option<&str>) -> Result<()> {
let n = sqlx::query(
"UPDATE users SET username = ?2, display_name = ?3, updated_at = datetime('now')
WHERE id = ?1",
)
.bind(id)
.bind(username)
.bind(display_name)
.execute(pool)
.await?
.rows_affected();
if n == 0 {
bail!("no such user: {id}");
}
Ok(())
}
/// Removes the directory row only. The caller still owns `database/{id}.db`:
/// erasing a user means deleting that file too.
pub async fn delete(pool: &SqlitePool, id: &str) -> Result<()> {
sqlx::query("DELETE FROM users WHERE id = ?1")
.bind(id)
.execute(pool)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
/// Nested on purpose: also covers `init_system_pool` creating the parent directory.
fn temp_db_path(tag: &str) -> String {
let mut p = std::env::temp_dir();
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
p.push(format!("skald-test-{tag}-{}-{nanos}", std::process::id()));
p.push("database");
p.push("system.db");
p.to_string_lossy().into_owned()
}
fn cleanup(path: &str) {
if let Some(dir) = std::path::Path::new(path).parent().and_then(|p| p.parent()) {
let _ = std::fs::remove_dir_all(dir);
}
}
fn encrypted() -> Credentials {
Credentials::Encrypted {
kdf_params: r#"{"algo":"argon2id","m":65536,"t":3,"p":1}"#.into(),
kdf_salt: vec![1, 2, 3, 4],
database_password: vec![0xDE, 0xAD, 0xBE, 0xEF],
}
}
fn cleartext() -> Credentials {
Credentials::Cleartext(Some(ClearVerifier {
kdf_params: r#"{"algo":"argon2id","m":65536,"t":3,"p":1}"#.into(),
kdf_salt: vec![5, 6, 7, 8],
password_hash: vec![0xAB, 0xCD],
}))
}
#[tokio::test]
async fn init_system_pool_creates_the_database_directory() {
let path = temp_db_path("users-mkdir");
assert!(!std::path::Path::new(&path).parent().unwrap().exists());
let pool = crate::db::init_system_pool(&path).await.unwrap();
assert!(std::path::Path::new(&path).exists(), "system.db must exist under a fresh database/");
pool.close().await;
cleanup(&path);
}
#[tokio::test]
async fn encrypted_user_round_trips_and_keeps_the_wrapped_dek() {
let path = temp_db_path("users-enc");
let pool = crate::db::init_system_pool(&path).await.unwrap();
insert(&pool, "u-1", "ada", Some("Ada"), "admin", &encrypted()).await.unwrap();
let u = by_username(&pool, "ada").await.unwrap().expect("user by username");
assert_eq!(u.id, "u-1");
assert!(u.is_encrypted());
assert!(u.active);
match u.credentials {
Credentials::Encrypted { database_password, kdf_salt, .. } => {
assert_eq!(database_password, vec![0xDE, 0xAD, 0xBE, 0xEF]);
assert_eq!(kdf_salt, vec![1, 2, 3, 4]);
}
other => panic!("expected Encrypted, got {other:?}"),
}
pool.close().await;
cleanup(&path);
}
#[tokio::test]
async fn cleartext_user_round_trips_with_and_without_a_verifier() {
let path = temp_db_path("users-clear");
let pool = crate::db::init_system_pool(&path).await.unwrap();
insert(&pool, "u-1", "kid", None, "children", &cleartext()).await.unwrap();
insert(&pool, "u-2", "kiosk", None, "children", &Credentials::Cleartext(None)).await.unwrap();
let with = get(&pool, "u-1").await.unwrap().unwrap();
assert!(!with.is_encrypted());
match with.credentials {
Credentials::Cleartext(Some(v)) => assert_eq!(v.password_hash, vec![0xAB, 0xCD]),
other => panic!("expected a verifier, got {other:?}"),
}
let without = get(&pool, "u-2").await.unwrap().unwrap();
assert!(matches!(without.credentials, Credentials::Cleartext(None)));
assert_eq!(count(&pool).await.unwrap(), 2);
assert_eq!(list(&pool).await.unwrap().len(), 2);
pool.close().await;
cleanup(&path);
}
/// Changing the password re-wraps the DEK; migrating to cleartext must clear it.
#[tokio::test]
async fn set_credentials_rewraps_and_migrates() {
let path = temp_db_path("users-rewrap");
let pool = crate::db::init_system_pool(&path).await.unwrap();
insert(&pool, "u-1", "ada", None, "admin", &encrypted()).await.unwrap();
let rewrapped = Credentials::Encrypted {
kdf_params: r#"{"algo":"argon2id","m":65536,"t":3,"p":1}"#.into(),
kdf_salt: vec![9, 9, 9],
database_password: vec![0xFE, 0xED],
};
set_credentials(&pool, "u-1", &rewrapped).await.unwrap();
match get(&pool, "u-1").await.unwrap().unwrap().credentials {
Credentials::Encrypted { database_password, .. } => assert_eq!(database_password, vec![0xFE, 0xED]),
other => panic!("expected Encrypted, got {other:?}"),
}
set_credentials(&pool, "u-1", &cleartext()).await.unwrap();
let u = get(&pool, "u-1").await.unwrap().unwrap();
assert!(!u.is_encrypted(), "migrating must flip `encrypted` and drop the wrapped DEK");
assert!(set_credentials(&pool, "ghost", &cleartext()).await.is_err(), "unknown id must fail");
pool.close().await;
cleanup(&path);
}
/// The SQL `CHECK` is the last line of defence when a row is written without
/// going through [`Credentials`].
#[tokio::test]
async fn check_constraint_rejects_impossible_rows() {
let path = temp_db_path("users-check");
let pool = crate::db::init_system_pool(&path).await.unwrap();
// encrypted without a wrapped DEK
let err = sqlx::query(
"INSERT INTO users (id, username, role_id, encrypted) VALUES ('x', 'x', 'admin', 1)",
)
.execute(&pool)
.await;
assert!(err.is_err(), "encrypted=1 requires database_password");
// encrypted *and* carrying a password hash
let err = sqlx::query(
"INSERT INTO users (id, username, role_id, encrypted, database_password, password_hash)
VALUES ('y', 'y', 'admin', 1, X'00', X'01')",
)
.execute(&pool)
.await;
assert!(err.is_err(), "an encrypted user must not also store a password hash");
// cleartext carrying a wrapped DEK
let err = sqlx::query(
"INSERT INTO users (id, username, role_id, encrypted, database_password)
VALUES ('z', 'z', 'admin', 0, X'00')",
)
.execute(&pool)
.await;
assert!(err.is_err(), "cleartext=0 must not store a wrapped DEK");
assert_eq!(count(&pool).await.unwrap(), 0);
pool.close().await;
cleanup(&path);
}
#[tokio::test]
async fn debug_never_prints_key_material() {
let u = User {
id: "u-1".into(),
username: "ada".into(),
display_name: None,
role_id: "admin".into(),
credentials: encrypted(),
active: true,
created_at: "now".into(),
updated_at: "now".into(),
};
let printed = format!("{u:?}");
assert!(printed.contains("ada"));
assert!(!printed.contains("222"), "no raw DEK bytes");
assert!(!printed.contains("deadbeef") && !printed.contains("DEADBEEF"));
assert!(printed.contains("<redacted>"));
}
}
+266
View File
@@ -0,0 +1,266 @@
//! Elicitation — server-initiated input requests (MCP spec 2025-06-18).
//!
//! When an MCP server needs input *during* a tool call (e.g. a sudo password),
//! it sends `elicitation/create`. The `mcp-client` read-loop forwards it through
//! the [`ElicitationHandler`] bridge to the [`ElicitationManager`], which surfaces
//! it in the Agent Inbox and waits for the user's decision. The reply (and any
//! secret it carries) flows straight back to the server's stdin — it is **never**
//! logged, broadcast in an event, or written to the DB.
//!
//! Mirrors [`crate::clarification`], but with the `accept`/`decline`/`cancel`
//! outcome and a `sensitive` flag that elicitation needs and clarification lacks.
use std::sync::Arc;
use std::sync::atomic::{AtomicI64, Ordering};
use std::time::Duration;
use async_trait::async_trait;
use chrono::Utc;
use serde::Serialize;
use serde_json::Value;
use tokio::sync::{broadcast, oneshot};
use tracing::{debug, info};
use mcp_client::{ElicitationAction, ElicitationHandler, ElicitationReply, ElicitationRequest};
use crate::events::{GlobalEvent, ServerEvent};
use crate::pending_registry::PendingRegistry;
/// How long the user has to answer an elicitation before we reply `cancel`.
/// Independent of any secret-cache TTL the MCP server keeps in its own RAM.
const ELICITATION_DEADLINE: Duration = Duration::from_secs(300);
/// One pending elicitation, surfaced to the Inbox UI. Holds **no value** — only
/// the prompt metadata. The secret travels through the `oneshot`, not here.
#[derive(Debug, Clone, Serialize)]
pub struct PendingElicitationInfo {
pub request_id: i64,
pub server_name: String,
pub message: String,
/// Name of the single requested field (v1 supports one field), if any.
pub field_name: Option<String>,
/// Render the input masked (`<input type="password">`) and never echo it.
pub sensitive: bool,
/// Empty `requestedSchema` ⇒ pure yes/no confirmation (no input field).
pub is_confirmation: bool,
pub created_at: String,
}
/// The user's decision, fed back from the Inbox API into the waiting handler.
#[derive(Debug, Clone)]
pub struct ElicitationOutcome {
/// `"accept"` | `"decline"` | `"cancel"`.
pub action: String,
/// Field values for `accept` (e.g. `{ "password": "…" }`); `None` otherwise.
pub content: Option<Value>,
}
pub struct ElicitationManager {
/// Shared pending-request plumbing (map + oneshot). Keyed by `request_id` —
/// elicitation is server-initiated and has no durable `tool_call_id`.
registry: PendingRegistry<PendingElicitationInfo, ElicitationOutcome>,
next_id: AtomicI64,
/// Global event bus, mirroring `ClarificationManager`. Broadcasts
/// `ElicitationRequested` / `ElicitationResolved` so Inbox subscribers
/// re-snapshot. **Never** carries the secret — only `request_id` + title.
event_tx: broadcast::Sender<GlobalEvent>,
}
impl ElicitationManager {
pub fn new(event_tx: broadcast::Sender<GlobalEvent>) -> Arc<Self> {
Arc::new(Self {
registry: PendingRegistry::new(),
next_id: AtomicI64::new(1),
event_tx,
})
}
/// Register a pending elicitation derived from an `elicitation/create`
/// request. Returns the id and a receiver that resolves when the user
/// answers (via the Inbox API) or the request is cancelled.
pub async fn register(
&self,
server_name: &str,
message: &str,
requested_schema: &Value,
) -> (i64, oneshot::Receiver<ElicitationOutcome>) {
let request_id = self.next_id.fetch_add(1, Ordering::SeqCst);
let (field_name, sensitive, is_confirmation) = parse_schema(requested_schema);
let info = PendingElicitationInfo {
request_id,
server_name: server_name.to_string(),
message: message.to_string(),
field_name,
sensitive,
is_confirmation,
created_at: Utc::now().to_rfc3339(),
};
let title = if message.is_empty() {
format!("{server_name}: input requested")
} else {
message.to_string()
};
let rx = self.registry.insert(request_id, info).await;
info!(server = server_name, request_id, sensitive, "elicitation: pending registered");
let _ = self.event_tx.send(GlobalEvent {
source: None,
session_id: None,
event: ServerEvent::ElicitationRequested { request_id, title },
});
(request_id, rx)
}
/// Resolve a pending elicitation with the user's decision. The `content`
/// (which may hold a secret) is forwarded on the `oneshot` and never logged.
pub async fn resolve(&self, request_id: i64, outcome: ElicitationOutcome) -> bool {
let action = outcome.action.clone();
if self.registry.resolve(request_id, outcome).await.is_some() {
debug!(request_id, %action, "elicitation: resolved");
self.broadcast_resolved(request_id);
true
} else {
false
}
}
/// Drop a pending elicitation without a user answer (deadline elapsed or the
/// waiting handler went away). The dropped `oneshot` sender makes the handler
/// reply `cancel`.
pub async fn cancel(&self, request_id: i64) {
if self.registry.remove(request_id).await.is_some() {
debug!(request_id, "elicitation: cancelled (deadline/handler gone)");
self.broadcast_resolved(request_id);
}
}
pub async fn list_pending(&self) -> Vec<PendingElicitationInfo> {
let mut items = self.registry.list().await;
items.sort_by(|a, b| a.created_at.cmp(&b.created_at));
items
}
fn broadcast_resolved(&self, request_id: i64) {
let _ = self.event_tx.send(GlobalEvent {
source: None,
session_id: None,
event: ServerEvent::ElicitationResolved { request_id },
});
}
}
/// Derives, from an MCP `requestedSchema`, the single field name, whether it is
/// sensitive (masked input), and whether it is a pure confirmation (empty schema).
/// v1 supports exactly one field — extra properties are ignored.
fn parse_schema(schema: &Value) -> (Option<String>, bool, bool) {
match schema.get("properties").and_then(Value::as_object) {
Some(props) if !props.is_empty() => {
let (key, def) = props.iter().next().unwrap();
let format = def.get("format").and_then(Value::as_str).unwrap_or("");
let write_only = def.get("writeOnly").and_then(Value::as_bool).unwrap_or(false);
let name_l = key.to_lowercase();
let sensitive = format == "password"
|| write_only
|| ["password", "passphrase", "secret", "token"]
.iter()
.any(|s| name_l.contains(s));
(Some(key.clone()), sensitive, false)
}
// No fields ⇒ confirmation request.
_ => (None, false, true),
}
}
/// Bridges `mcp-client`'s server→client elicitation to the `ElicitationManager`.
/// Registers the request, waits up to [`ELICITATION_DEADLINE`] for the user, and
/// maps the outcome back to an [`ElicitationReply`].
pub struct ElicitationBridge {
manager: Arc<ElicitationManager>,
}
impl ElicitationBridge {
pub fn new(manager: Arc<ElicitationManager>) -> Arc<Self> {
Arc::new(Self { manager })
}
}
#[async_trait]
impl ElicitationHandler for ElicitationBridge {
async fn handle(&self, server_name: &str, request: ElicitationRequest) -> ElicitationReply {
let (id, rx) = self
.manager
.register(server_name, &request.message, &request.requested_schema)
.await;
match tokio::time::timeout(ELICITATION_DEADLINE, rx).await {
Ok(Ok(outcome)) => {
let action = match outcome.action.as_str() {
"accept" => ElicitationAction::Accept,
"decline" => ElicitationAction::Decline,
_ => ElicitationAction::Cancel,
};
let content = if action == ElicitationAction::Accept { outcome.content } else { None };
ElicitationReply { action, content }
}
// Deadline elapsed or the resolver's sender was dropped → cancel.
_ => {
self.manager.cancel(id).await;
ElicitationReply { action: ElicitationAction::Cancel, content: None }
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn empty_schema_is_confirmation() {
let (field, sensitive, confirm) = parse_schema(&json!({ "type": "object", "properties": {} }));
assert_eq!(field, None);
assert!(!sensitive);
assert!(confirm);
}
#[test]
fn missing_properties_is_confirmation() {
let (field, _sensitive, confirm) = parse_schema(&json!({ "type": "object" }));
assert_eq!(field, None);
assert!(confirm);
}
#[test]
fn password_format_is_sensitive() {
let schema = json!({ "type": "object", "properties": {
"password": { "type": "string", "format": "password" }
}});
let (field, sensitive, confirm) = parse_schema(&schema);
assert_eq!(field.as_deref(), Some("password"));
assert!(sensitive);
assert!(!confirm);
}
#[test]
fn secret_by_name_is_sensitive() {
let schema = json!({ "type": "object", "properties": {
"api_token": { "type": "string" }
}});
let (_field, sensitive, _confirm) = parse_schema(&schema);
assert!(sensitive);
}
#[test]
fn plain_field_is_not_sensitive() {
let schema = json!({ "type": "object", "properties": {
"hostname": { "type": "string" }
}});
let (field, sensitive, confirm) = parse_schema(&schema);
assert_eq!(field.as_deref(), Some("hostname"));
assert!(!sensitive);
assert!(!confirm);
}
}
+3
View File
@@ -0,0 +1,3 @@
pub use core_api::events::{
ClientMessage, GlobalEvent, InboundDataMessage, ServerEvent,
};
+103
View File
@@ -0,0 +1,103 @@
use anyhow::{Context, Result};
use sqlx::SqlitePool;
use super::ImageGenerateModelRecord;
#[derive(sqlx::FromRow)]
struct ImageGenerateModelRow {
id: i64,
provider_id: i64,
model_id: String,
name: String,
priority: i64,
}
pub async fn load_all(pool: &SqlitePool) -> Result<Vec<ImageGenerateModelRecord>> {
let rows = sqlx::query_as::<_, ImageGenerateModelRow>(
"SELECT id, provider_id, model_id, name, priority
FROM image_generate_models
WHERE removed_at IS NULL
ORDER BY priority ASC, name ASC",
)
.fetch_all(pool)
.await
.context("image_generate_models: load_all")?;
Ok(rows.into_iter().map(row_to_record).collect())
}
pub async fn insert(pool: &SqlitePool, r: &ImageGenerateModelRecord) -> Result<i64> {
let restored = sqlx::query_scalar::<_, i64>(
"UPDATE image_generate_models
SET provider_id=?1, model_id=?2, name=?3, priority=?4, removed_at=NULL
WHERE id = (
SELECT id FROM image_generate_models
WHERE removed_at IS NOT NULL
AND (provider_id=?1 AND model_id=?2 OR name=?3)
LIMIT 1
)
RETURNING id",
)
.bind(r.provider_id)
.bind(&r.model_id)
.bind(&r.name)
.bind(r.priority as i64)
.fetch_optional(pool)
.await
.context("image_generate_models: restore soft-deleted")?;
if let Some(id) = restored {
return Ok(id);
}
sqlx::query_scalar::<_, i64>(
"INSERT INTO image_generate_models (provider_id, model_id, name, priority)
VALUES (?1, ?2, ?3, ?4)
RETURNING id",
)
.bind(r.provider_id)
.bind(&r.model_id)
.bind(&r.name)
.bind(r.priority as i64)
.fetch_one(pool)
.await
.context("image_generate_models: insert")
}
pub async fn update(pool: &SqlitePool, id: i64, r: &ImageGenerateModelRecord) -> Result<()> {
sqlx::query(
"UPDATE image_generate_models
SET provider_id=?1, model_id=?2, name=?3, priority=?4
WHERE id=?5",
)
.bind(r.provider_id)
.bind(&r.model_id)
.bind(&r.name)
.bind(r.priority as i64)
.bind(id)
.execute(pool)
.await
.context("image_generate_models: update")?;
Ok(())
}
pub async fn soft_delete(pool: &SqlitePool, id: i64) -> Result<()> {
sqlx::query(
"UPDATE image_generate_models SET removed_at = datetime('now') WHERE id = ?1",
)
.bind(id)
.execute(pool)
.await
.context("image_generate_models: soft-delete")?;
Ok(())
}
fn row_to_record(r: ImageGenerateModelRow) -> ImageGenerateModelRecord {
ImageGenerateModelRecord {
id: r.id,
provider_id: r.provider_id,
model_id: r.model_id,
name: r.name,
priority: r.priority as i32,
}
}
@@ -0,0 +1,296 @@
/// ImageGeneratorManager — DB-aware registry of image generation providers.
///
/// Two kinds of providers coexist:
/// - **DB-backed**: rows in `image_generate_models`, built from `llm_providers` credentials.
/// Managed via `add_model` / `update_model` / `delete_model`. Loaded on startup
/// and after every mutation.
/// - **Plugin-registered**: ephemeral providers registered at runtime by plugins.
/// Not persisted — they disappear on plugin stop.
///
/// `get(id)` resolves by explicit id across both plugin and DB-backed providers.
/// When called without an id, plugin providers take precedence over DB-backed ones.
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::{Context, Result, anyhow};
use async_trait::async_trait;
use rand::RngExt;
use sqlx::SqlitePool;
use tokio::sync::RwLock;
use tracing::{info, warn};
use core_api::image_generate::ImageGenerateRegistry;
use crate::llm::LlmProviderRecord;
use crate::llm::db as llm_db;
use crate::provider::ProviderRegistry;
use crate::tools::Tool;
use super::{ImageGenerate, ImageGenerateInfo, ImageGenerateModelInfo, ImageGenerateModelRecord};
use super::db as image_db;
// ── Internal state ────────────────────────────────────────────────────────────
struct ImageGenerateSlot {
record: ImageGenerateModelRecord,
provider: LlmProviderRecord,
generator: Arc<dyn ImageGenerate>,
}
struct ManagerState {
/// DB-backed generators, ordered by priority ASC. Rebuilt on every reload().
db_slots: Vec<ImageGenerateSlot>,
/// Plugin-registered providers (ephemeral — not in DB).
plugins: Vec<Arc<dyn ImageGenerate>>,
}
// ── ImageGeneratorManager ─────────────────────────────────────────────────────
pub struct ImageGeneratorManager {
pool: Arc<SqlitePool>,
registry: Arc<ProviderRegistry>,
state: RwLock<ManagerState>,
data_root: PathBuf,
}
impl ImageGeneratorManager {
pub async fn new(
pool: Arc<SqlitePool>,
registry: Arc<ProviderRegistry>,
data_root: impl Into<PathBuf>,
) -> Result<Arc<Self>> {
let mgr = Arc::new(Self {
pool,
registry,
state: RwLock::new(ManagerState {
db_slots: Vec::new(),
plugins: Vec::new(),
}),
data_root: data_root.into(),
});
mgr.reload().await?;
Ok(mgr)
}
// ── Plugin registration (ephemeral) ───────────────────────────────────────
pub async fn register(&self, provider: Arc<dyn ImageGenerate>) {
let mut state = self.state.write().await;
let id = provider.id().to_string();
state.plugins.retain(|p| p.id() != id);
state.plugins.push(provider);
info!(provider_id = %id, "image generator registered (plugin)");
}
pub async fn unregister(&self, id: &str) {
let mut state = self.state.write().await;
let before = state.plugins.len();
state.plugins.retain(|p| p.id() != id);
if state.plugins.len() < before {
info!(provider_id = %id, "image generator unregistered (plugin)");
}
}
// ── Model CRUD (DB-backed) ────────────────────────────────────────────────
pub async fn add_model(&self, record: ImageGenerateModelRecord) -> Result<i64> {
let id = image_db::insert(&self.pool, &record).await?;
self.reload().await?;
Ok(id)
}
pub async fn update_model(&self, id: i64, record: ImageGenerateModelRecord) -> Result<()> {
image_db::update(&self.pool, id, &record).await?;
self.reload().await
}
pub async fn delete_model(&self, id: i64) -> Result<()> {
image_db::soft_delete(&self.pool, id).await?;
self.reload().await
}
pub async fn get_model(&self, id: i64) -> Option<ImageGenerateModelRecord> {
self.state.read().await
.db_slots.iter()
.find(|s| s.record.id == id)
.map(|s| s.record.clone())
}
pub async fn list_models_info(&self) -> Vec<ImageGenerateModelInfo> {
self.state.read().await.db_slots.iter().map(|s| ImageGenerateModelInfo {
id: s.record.id,
provider_id: s.provider.id,
provider_name: s.provider.name.clone(),
model_id: s.record.model_id.clone(),
name: s.record.name.clone(),
priority: s.record.priority,
from_plugin: false,
description: None,
}).collect()
}
/// Returns all active providers: plugin-registered first, then DB-backed by priority.
pub async fn list_all_info(&self) -> Vec<ImageGenerateModelInfo> {
let state = self.state.read().await;
let plugins = state.plugins.iter().map(|p| ImageGenerateModelInfo {
id: 0,
provider_id: 0,
provider_name: "Plugin".into(),
model_id: p.id().to_string(),
name: p.name().to_string(),
priority: 0,
from_plugin: true,
description: p.description().map(str::to_string),
});
let db = state.db_slots.iter().map(|s| ImageGenerateModelInfo {
id: s.record.id,
provider_id: s.provider.id,
provider_name: s.provider.name.clone(),
model_id: s.record.model_id.clone(),
name: s.record.name.clone(),
priority: s.record.priority,
from_plugin: false,
description: None,
});
plugins.chain(db).collect()
}
// ── Provider queries ───────────────────────────────────────────────────────
/// Returns all active providers as lightweight info structs (for LLM tool).
pub async fn list(&self) -> Vec<ImageGenerateInfo> {
let state = self.state.read().await;
state.plugins.iter()
.map(|p| ImageGenerateInfo {
id: p.id().to_string(),
name: p.name().to_string(),
description: p.description().map(str::to_string),
extra_params_schema: p.extra_params_schema(),
})
.chain(state.db_slots.iter().map(|s| ImageGenerateInfo {
id: s.record.name.clone(),
name: s.record.name.clone(),
description: None,
extra_params_schema: None,
}))
.collect()
}
/// Looks up a provider by id — plugins first, then DB-backed by name.
pub async fn get(&self, id: &str) -> Option<Arc<dyn ImageGenerate>> {
let state = self.state.read().await;
if let Some(p) = state.plugins.iter().find(|p| p.id() == id) {
return Some(Arc::clone(p));
}
state.db_slots.iter()
.find(|s| s.record.name == id)
.map(|s| Arc::clone(&s.generator))
}
// ── Generation ────────────────────────────────────────────────────────────
pub async fn generate(
&self,
provider_id: &str,
prompt: &str,
extra_params: Option<&serde_json::Value>,
) -> Result<(PathBuf, String)> {
let provider = self.get(provider_id).await
.ok_or_else(|| anyhow!("image provider '{}' not found", provider_id))?;
let images_dir = self.data_root.join("images");
tokio::fs::create_dir_all(&images_dir).await?;
let bytes = provider.generate(prompt, extra_params).await?;
let file_id: String = rand::rng()
.sample_iter(rand::distr::Alphanumeric)
.take(32)
.map(char::from)
.collect();
let path = images_dir.join(format!("{file_id}.png"));
tokio::fs::write(&path, &bytes).await?;
let url = format!("/api/images/{file_id}");
info!(provider_id, path = %path.display(), "image generated");
Ok((path, url))
}
// ── Tool injection ─────────────────────────────────────────────────────────
/// Returns the two image tools when at least one provider is active.
/// Called per-turn by the session handler to conditionally inject tools.
pub async fn tools(self: Arc<Self>) -> Vec<Arc<dyn Tool>> {
let state = self.state.read().await;
if state.plugins.is_empty() && state.db_slots.is_empty() {
return vec![];
}
drop(state);
vec![
Arc::new(crate::tools::image_generate::ImageGenerateProvidersList { mgr: Arc::clone(&self) }) as Arc<dyn Tool>,
Arc::new(crate::tools::image_generate::ImageGenerateTool { mgr: Arc::clone(&self) }) as Arc<dyn Tool>,
]
}
pub fn images_dir(&self) -> PathBuf {
self.data_root.join("images")
}
// ── Private ───────────────────────────────────────────────────────────────
async fn reload(&self) -> Result<()> {
let model_records = image_db::load_all(&self.pool).await?;
let provider_records: Vec<LlmProviderRecord> =
llm_db::load_all_providers(&self.pool).await?;
let providers: std::collections::HashMap<i64, LlmProviderRecord> =
provider_records.into_iter().map(|p| (p.id, p)).collect();
let mut db_slots = Vec::new();
for model in model_records {
let provider = match providers.get(&model.provider_id) {
Some(p) => p.clone(),
None => {
warn!(
model = %model.name,
provider_id = model.provider_id,
"orphaned image model — provider not found, skipping",
);
continue;
}
};
let result = self.registry.get(&provider.provider)
.and_then(|p| p.build_image_generator(&provider, &model))
.unwrap_or_else(|| anyhow::bail!("provider '{}' does not support image generation", provider.provider));
match result {
Ok(generator) => db_slots.push(ImageGenerateSlot { record: model, provider, generator }),
Err(e) => warn!(model = %model.name, error = %e, "failed to build image generator, skipping"),
}
}
let slot_count = db_slots.len();
self.state.write().await.db_slots = db_slots;
info!(db_backed = slot_count, "image generator manager reloaded");
Ok(())
}
}
// ── ImageGenerateRegistry impl ────────────────────────────────────────────────
#[async_trait]
impl ImageGenerateRegistry for ImageGeneratorManager {
async fn register(&self, provider: Arc<dyn ImageGenerate>) {
ImageGeneratorManager::register(self, provider).await;
}
async fn unregister(&self, id: &str) {
ImageGeneratorManager::unregister(self, id).await;
}
}
@@ -0,0 +1,37 @@
mod db;
pub mod manager;
pub mod openrouter_image;
pub use core_api::image_generate::ImageGenerate;
pub use core_api::image_generate::ImageGenerateModelRecord;
pub use manager::ImageGeneratorManager;
/// Public model metadata for API responses.
#[derive(Debug, Clone, serde::Serialize)]
pub struct ImageGenerateModelInfo {
pub id: i64,
pub provider_id: i64,
pub provider_name: String,
pub model_id: String,
pub name: String,
pub priority: i32,
/// `true` for plugin-registered (ephemeral) providers — not editable via the UI.
pub from_plugin: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
}
// ── Tool-facing types ─────────────────────────────────────────────────────────
/// Lightweight provider listing returned by `image_generate_providers_list`.
#[derive(Debug, Clone, serde::Serialize)]
pub struct ImageGenerateInfo {
pub id: String,
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
/// JSON Schema for the `extra_params` argument. Present only if the provider
/// accepts provider-specific parameters (e.g. width, height, steps).
#[serde(skip_serializing_if = "Option::is_none")]
pub extra_params_schema: Option<serde_json::Value>,
}
@@ -0,0 +1,94 @@
/// OpenRouter image generation via the chat completions endpoint with `modalities`.
///
/// Calls `POST {base_url}/chat/completions` with:
/// `{"model": ..., "messages": [...], "modalities": ["image", "text"]}`
///
/// The response image is returned as a base64 data URL inside
/// `choices[0].message.images[0].image_url.url`.
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use base64::Engine;
use tracing::{debug, info};
use super::ImageGenerate;
pub struct OpenRouterImageGenerator {
/// Stable display identifier, e.g. `"my_openrouter_grok"`.
id: String,
base_url: String,
api_key: String,
model: String,
http: reqwest::Client,
}
impl OpenRouterImageGenerator {
pub fn new(
id: impl Into<String>,
base_url: impl Into<String>,
api_key: impl Into<String>,
model: impl Into<String>,
) -> Self {
Self {
id: id.into(),
base_url: base_url.into(),
api_key: api_key.into(),
model: model.into(),
http: reqwest::Client::new(),
}
}
}
#[async_trait]
impl ImageGenerate for OpenRouterImageGenerator {
fn id(&self) -> &str { &self.id }
fn name(&self) -> &str { &self.id }
async fn generate(&self, prompt: &str, _extra_params: Option<&serde_json::Value>) -> Result<Vec<u8>> {
debug!(model = %self.model, "openrouter_image: generating");
let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
let body = serde_json::json!({
"model": self.model,
"messages": [{ "role": "user", "content": prompt }],
"modalities": ["image"],
});
let resp = self.http
.post(&url)
.bearer_auth(&self.api_key)
.header("X-Title", core_api::APP_NAME)
.json(&body)
.send()
.await
.map_err(|e| anyhow!("openrouter_image: request failed: {e}"))?;
let status = resp.status();
let json: serde_json::Value = resp
.json()
.await
.map_err(|e| anyhow!("openrouter_image: response parse failed: {e}"))?;
if !status.is_success() {
let msg = json["error"]["message"].as_str().unwrap_or("unknown error");
anyhow::bail!("openrouter_image: API error {status}: {msg}");
}
let data_url = json["choices"][0]["message"]["images"][0]["image_url"]["url"]
.as_str()
.ok_or_else(|| anyhow!("openrouter_image: no image in response — full response: {json}"))?;
let b64 = data_url
.strip_prefix("data:image/png;base64,")
.or_else(|| data_url.strip_prefix("data:image/jpeg;base64,"))
.or_else(|| data_url.strip_prefix("data:image/webp;base64,"))
.unwrap_or(data_url);
let bytes = base64::engine::general_purpose::STANDARD
.decode(b64)
.map_err(|e| anyhow!("openrouter_image: base64 decode failed: {e}"))?;
info!(model = %self.model, bytes = bytes.len(), "openrouter_image: generation complete");
Ok(bytes)
}
}
+142
View File
@@ -0,0 +1,142 @@
use std::sync::Arc;
use async_trait::async_trait;
use serde::Serialize;
use serde_json::Value;
use core_api::inbox::{
InboxApi, InboxApprovalItem, InboxClarificationItem, InboxElicitationItem, InboxSnapshot,
};
use core_api::tool::ToolDescriptionLength;
use crate::approval::{ApprovalManager, PendingApprovalInfo};
use crate::clarification::{ClarificationManager, PendingClarificationInfo};
use crate::elicitation::{ElicitationManager, ElicitationOutcome, PendingElicitationInfo};
use crate::tools::ToolRegistry;
#[derive(Serialize)]
pub struct InboxItems {
pub total: usize,
pub approvals: Vec<PendingApprovalInfo>,
pub clarifications: Vec<PendingClarificationInfo>,
pub elicitations: Vec<PendingElicitationInfo>,
}
#[derive(Clone)]
pub struct Inbox {
pub approval: Arc<ApprovalManager>,
clarification: Arc<ClarificationManager>,
elicitation: Arc<ElicitationManager>,
/// Used to humanise approval tool calls (`describe`) when building snapshots.
tools: Arc<ToolRegistry>,
}
impl Inbox {
pub fn new(
approval: Arc<ApprovalManager>,
clarification: Arc<ClarificationManager>,
elicitation: Arc<ElicitationManager>,
tools: Arc<ToolRegistry>,
) -> Self {
Self { approval, clarification, elicitation, tools }
}
pub async fn list_pending(&self) -> InboxItems {
let mut approvals = self.approval.list_pending().await;
// Union in DB-persisted pending approvals not represented in memory, so the
// Inbox survives a server restart (the registry is in-memory only). Both sources
// key on the durable `tool_call_id` (live approvals now carry
// `request_id == tool_call_id`; persisted ones carry the falsy
// `PERSISTED_REQUEST_ID`, telling the client to resolve by `tool_call_id`), so the
// dedup below is a single-id-space set difference.
let live: std::collections::HashSet<i64> =
approvals.iter().map(|a| a.tool_call_id).collect();
for a in self.approval.list_persisted_pending().await {
if !live.contains(&a.tool_call_id) {
approvals.push(a);
}
}
let clarifications = self.clarification.list_pending().await;
let elicitations = self.elicitation.list_pending().await;
let total = approvals.len() + clarifications.len() + elicitations.len();
InboxItems { total, approvals, clarifications, elicitations }
}
pub async fn approve(&self, request_id: i64) {
self.approval.approve(request_id).await;
}
pub async fn reject(&self, request_id: i64, note: String) {
self.approval.reject(request_id, note).await;
}
pub async fn answer(&self, request_id: i64, answer: String) -> bool {
self.clarification.resolve(request_id, answer).await
}
pub async fn resolve_elicitation(&self, request_id: i64, action: String, content: Option<Value>) -> bool {
self.elicitation.resolve(request_id, ElicitationOutcome { action, content }).await
}
}
/// Exposes the Inbox to plugins via `PluginContext` (plugin.md §12.2). Converts
/// the main-crate pending types into the core-api snapshot types.
#[async_trait]
impl InboxApi for Inbox {
async fn list_pending(&self) -> InboxSnapshot {
let items = self.list_pending().await;
let approvals = items.approvals.into_iter().map(|a| {
// Humanise the tool call for the card / notification; ship the raw
// arguments untruncated so the detail dialog shows exactly what is
// being approved (e.g. the full `execute_cmd` command).
let summary = self.tools.describe_call(&a.tool_name, &a.arguments, ToolDescriptionLength::Short);
InboxApprovalItem {
request_id: a.request_id,
tool_name: a.tool_name,
summary,
arguments: a.arguments,
agent_id: a.agent_id,
source: a.source,
context_label: a.context_label,
created_at: a.created_at,
}
}).collect();
let clarifications = items.clarifications.into_iter().map(|c| InboxClarificationItem {
request_id: c.request_id,
agent_id: c.agent_id,
source: c.source,
context_label: c.context_label,
title: c.title,
question: c.question,
suggested_answers: c.suggested_answers,
created_at: c.created_at,
}).collect();
let elicitations = items.elicitations.into_iter().map(|e| InboxElicitationItem {
request_id: e.request_id,
server_name: e.server_name,
message: e.message,
field_name: e.field_name,
sensitive: e.sensitive,
is_confirmation: e.is_confirmation,
created_at: e.created_at,
}).collect();
InboxSnapshot { total: items.total, approvals, clarifications, elicitations }
}
async fn approve(&self, request_id: i64) {
self.approve(request_id).await;
}
async fn reject(&self, request_id: i64, reason: String) {
self.reject(request_id, reason).await;
}
async fn answer(&self, request_id: i64, answer: String) -> bool {
self.answer(request_id, answer).await
}
async fn resolve_elicitation(&self, request_id: i64, action: String, content: Option<Value>) -> bool {
self.resolve_elicitation(request_id, action, content).await
}
}
+695
View File
@@ -0,0 +1,695 @@
//! `LatexCompiler` — compiles `.tex` sources to PDF using `latexmk -xelatex`.
//!
//! ## Caching (dependency-aware)
//! LaTeX documents routinely pull in external fragments via `\input`,
//! `\include`, `\includegraphics`, custom `.sty`/`.cls` packages, `.bib`
//! files, and so on. A cache keyed only on the main `.tex` content would serve
//! stale PDFs whenever one of those dependencies changes, so we use the
//! `.fls` recorder file produced by TeX (and orchestrated by `latexmk`) to
//! discover the full set of inputs and key the cache on their combined
//! content hash.
//!
//! Two cache artefacts live under `<tmp>/skald-latex/`:
//!
//! | Artefact | Key | Purpose |
//! |-----------------------|-------------------------------------|------------------------------------------|
//! | `<path-hash>.fls` | SHA-256 of the `.tex` absolute path | Last-known input list for that source |
//! | `<deps-hash>.pdf` | SHA-256 of every input's contents | The compiled PDF for that exact state |
//!
//! Lookup flow per request:
//! 1. Read `<path-hash>.fls`. If missing → fresh compile.
//! 2. Parse it, keep only user-controlled inputs (see [`parse_user_deps`]),
//! hash every file's bytes, derive `<deps-hash>`.
//! 3. If `<deps-hash>.pdf` exists → cache hit, serve it.
//! 4. Otherwise → run `latexmk`, capture the new `.fls`, overwrite the
//! `<path-hash>.fls` sidecar, save the PDF as `<deps-hash>.pdf`, serve.
//!
//! `latexmk` runs in a per-compile scratch directory (`-output-directory`)
//! using the source file's own directory as CWD, so relative
//! `\input`/`\includegraphics` references resolve as they would in a local
//! build. The scratch directory is removed before returning.
//!
//! ## Failure modes
//! - `ToolMissing` — `latexmk` is not on `PATH` (e.g. no TeX distribution).
//! - `Timeout` — compilation exceeded [`COMPILE_TIMEOUT_SECS`].
//! - `Failed { log }` — `latexmk` exited non-zero; the `.log` is captured so
//! callers can surface a useful message (the file viewer falls back to plain
//! text in this case).
//!
//! ## Residual limitations
//! - System TeX packages (under [`TEXMF_PREFIXES`]) are deliberately excluded
//! from the dependency hash — they only change with a TeX distribution
//! upgrade, which is rare and easy to handle by clearing the cache.
//! - Files consumed via `\input{|"shell command"}` (shell-escape) are not
//! recorded in the `.fls`; documents relying on this will not invalidate
//! the cache properly. Acceptable for V1.
use std::path::{Path, PathBuf};
use std::time::Duration;
use sha2::{Digest, Sha256};
use tokio::process::Command;
/// Hard ceiling for a single `latexmk` run. `latexmk` itself never prompts
/// under `-interaction=nonstopmode`, but packages can still hang (e.g. waiting
/// on missing fonts); the timeout guards against that.
const COMPILE_TIMEOUT_SECS: u64 = 30;
/// Default subdirectory of the OS temp dir used to store cached PDFs and
/// per-compile scratch directories.
const CACHE_DIR_NAME: &str = "skald-latex";
/// Extensions of files TeX produces as side-effects of compilation. They are
/// written to the output directory alongside the PDF and never count as
/// user-controlled dependencies.
const AUX_EXTS: &[&str] = &[
"aux", "log", "fls", "fdb_latexmk", "synctex.gz", "out",
"toc", "bbl", "blg", "run.xml", "idx", "ind", "ilg",
"lof", "lot", "nav", "snm", "vrb", "bcf", "xdv", "mtc",
];
/// Path prefixes that identify a system TeX distribution. Files matched here
/// (e.g. `/usr/local/texlive/2024/texmf-dist/.../article.cls`) are filtered out
/// of the dependency set: they only change on a distro upgrade, which is rare
/// and easy to handle by clearing the cache manually.
const TEXMF_PREFIXES: &[&str] = &[
"/usr/local/texlive",
"/Library/TeX",
"/opt/homebrew/texlive",
"/usr/share/texmf",
"/usr/share/texlive",
"/var/lib/texmf",
];
/// A successfully compiled PDF.
pub struct CompiledPdf {
pub bytes: Vec<u8>,
/// `true` when served from cache without invoking `latexmk`. Currently
/// informational only — surfaced in caller-side metrics/telemetry when
/// needed; kept on the struct so the API stays stable.
#[allow(dead_code)]
pub from_cache: bool,
}
/// Why a compilation request did not yield a PDF.
#[derive(Debug)]
pub enum CompileError {
/// `latexmk` is not reachable on `PATH`.
ToolMissing,
/// `latexmk` ran but exited with a non-zero status. Carries the textual
/// `.log` (or a synthetic message when the log is unavailable).
Failed { log: String },
/// Compilation did not finish within [`COMPILE_TIMEOUT_SECS`].
Timeout,
/// Underlying I/O error (reading the source, writing the cache, etc.).
Io(std::io::Error),
}
impl std::fmt::Display for CompileError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::ToolMissing => write!(f, "latexmk is not available on the server"),
Self::Failed { log } => write!(f, "compilation failed:\n{log}"),
Self::Timeout => write!(f, "compilation aborted (timeout {COMPILE_TIMEOUT_SECS}s)"),
Self::Io(e) => write!(f, "I/O error: {e}"),
}
}
}
impl std::error::Error for CompileError {}
impl From<std::io::Error> for CompileError {
fn from(e: std::io::Error) -> Self { Self::Io(e) }
}
/// Stateless-ish facade around `latexmk`. Owns only the cache root path; safe
/// to share via `Arc` (constructed once and stored on `Skald`).
#[derive(Clone)]
pub struct LatexCompiler {
cache_dir: PathBuf,
}
impl LatexCompiler {
pub fn new() -> Self {
Self { cache_dir: std::env::temp_dir().join(CACHE_DIR_NAME) }
}
/// Compile `tex_path` into a PDF, serving from cache when possible.
///
/// Cache lookup is **dependency-aware**: we first consult the `.fls`
/// sidecar that records every input TeX read on the last compile of this
/// source, hash the contents of those input files, and look up the PDF by
/// that composite hash. This means a change to any `\input`'ed fragment,
/// custom `.sty`, `.bib`, or `\includegraphics` target invalidates the
/// cache correctly even when the main `.tex` file is unchanged. See the
/// module docs for the full algorithm.
pub async fn compile(&self, tex_path: &Path) -> Result<CompiledPdf, CompileError> {
let path_key = path_hash(tex_path);
let fls_sidecar = self.cache_dir.join(format!("{path_key}.fls"));
// ── Cache lookup ───────────────────────────────────────────────────
// Read the cached .fls from the last compile of this exact path; if
// present, derive the composite deps hash and look for the PDF.
if let Ok(fls_text) = tokio::fs::read_to_string(&fls_sidecar).await {
let deps = parse_user_deps(&fls_text, tex_path);
match composite_hash_of(&deps).await {
Ok(deps_key) => {
let cached_pdf = self.cache_dir.join(format!("{deps_key}.pdf"));
if let Ok(bytes) = tokio::fs::read(&cached_pdf).await {
tracing::debug!(
?cached_pdf, deps_count = deps.len(),
"latex cache hit (deps-aware)"
);
return Ok(CompiledPdf { bytes, from_cache: true });
}
}
Err(e) => {
// One of the recorded deps is missing/unreadable — most
// likely a `\input` was deleted. Treat as a miss and
// recompile, which will refresh the `.fls` sidecar.
tracing::debug!(
error = %e, sidecar = ?fls_sidecar,
"deps hashing failed — falling through to fresh compile"
);
}
}
}
// ── Cache miss → compile ───────────────────────────────────────────
let (pdf_bytes, fresh_fls) = self.fresh_compile(tex_path, &path_key).await?;
// Persist the new .fls sidecar (overwrites the previous one for this
// path). A write failure is non-fatal: the next request will simply
// recompile again.
if let Err(e) = tokio::fs::write(&fls_sidecar, &fresh_fls).await {
tracing::warn!(?fls_sidecar, error = %e, "fls sidecar write failed");
}
// Compute the composite hash from the freshly recorded deps and store
// the PDF under that key.
let deps = parse_user_deps(&fresh_fls, tex_path);
let deps_key = composite_hash_of(&deps)
.await
.unwrap_or_else(|_| path_key.clone()); // fallback: at least cache by path
let cached_pdf = self.cache_dir.join(format!("{deps_key}.pdf"));
if let Err(e) = tokio::fs::write(&cached_pdf, &pdf_bytes).await {
tracing::warn!(?cached_pdf, error = %e, "latex cache write failed");
}
tracing::info!(
file = ?tex_path, deps_count = deps.len(),
"latex compiled (cache miss)"
);
Ok(CompiledPdf { bytes: pdf_bytes, from_cache: false })
}
/// Paths that should be watched to detect any change affecting the compiled
/// output of `tex_path`. Returns the source file itself plus every
/// user-controlled dependency listed in the cached `.fls` sidecar (the
/// recorder file from the last compile).
///
/// Returns just `[tex_path]` when no `.fls` is cached yet (e.g. before the
/// first compile), so the file watcher can install at least a baseline
/// watcher — once the first compile happens and the `.fls` is written, the
/// caller can call this again to pick up the full dependency set.
///
/// This is a synchronous best-effort read: a missing or unreadable `.fls`
/// is treated as "no deps known" rather than an error.
pub fn watch_paths_for(&self, tex_path: &Path) -> Vec<PathBuf> {
let mut paths = vec![tex_path.to_path_buf()];
let path_key = path_hash(tex_path);
let fls_sidecar = self.cache_dir.join(format!("{path_key}.fls"));
if let Ok(fls_text) = std::fs::read_to_string(&fls_sidecar) {
for dep in parse_user_deps(&fls_text, tex_path) {
if !paths.contains(&dep) {
paths.push(dep);
}
}
}
paths
}
/// Run `latexmk` for `tex_path` and return both the produced PDF bytes
/// and the textual `.fls` recorder file.
///
/// Uses a per-hash scratch directory under [`CACHE_DIR_NAME`] as the
/// `-output-directory`, while keeping the source file's directory as CWD so
/// relative `\input`/`\includegraphics` references resolve normally. The
/// scratch directory is removed before returning, regardless of outcome.
///
/// `path_key` is used only to namespace the scratch directory; it does not
/// affect the produced artefacts.
async fn fresh_compile(
&self,
tex_path: &Path,
path_key: &str,
) -> Result<(Vec<u8>, String), CompileError> {
if find_on_path("latexmk").await.is_none() {
return Err(CompileError::ToolMissing);
}
// Use a unique suffix so concurrent compiles of the same source (e.g.
// two requests racing before the .fls sidecar is written) do not
// collide on the scratch directory.
let out_dir = self.cache_dir.join(format!("{path_key}-{}/", unique_suffix()));
tokio::fs::create_dir_all(&out_dir).await?;
// CWD = source file's directory; falls back to "." for unusual inputs.
let cwd = tex_path.parent().unwrap_or_else(|| Path::new("."));
let mut cmd = Command::new("latexmk");
cmd.args([
"-xelatex",
"-interaction=nonstopmode",
"-halt-on-error",
"-file-line-error",
"-recorder", // ensure .fls is always produced
]);
cmd.arg(format!("-output-directory={}", out_dir.display()));
cmd.arg(tex_path);
cmd.current_dir(cwd);
cmd.stdout(std::process::Stdio::piped());
cmd.stderr(std::process::Stdio::piped());
// If our future is dropped (e.g. on shutdown) ensure the process dies.
cmd.kill_on_drop(true);
let output = match tokio::time::timeout(
Duration::from_secs(COMPILE_TIMEOUT_SECS),
cmd.output(),
).await {
Ok(Ok(o)) => o,
Ok(Err(e)) => {
let _ = cleanup_dir(&out_dir).await;
return Err(CompileError::Io(e));
}
Err(_) => {
// Timeout: `cmd.output()` future is dropped here; `kill_on_drop`
// takes care of terminating `latexmk`.
let _ = cleanup_dir(&out_dir).await;
return Err(CompileError::Timeout);
}
};
if !output.status.success() {
let log = read_compile_log(&out_dir, tex_path).await;
let _ = cleanup_dir(&out_dir).await;
return Err(CompileError::Failed { log });
}
let stem = file_stem(tex_path).unwrap_or_else(|| "output".to_string());
let pdf_path = out_dir.join(format!("{stem}.pdf"));
let fls_path = out_dir.join(format!("{stem}.fls"));
let pdf_bytes = match tokio::fs::read(&pdf_path).await {
Ok(b) => b,
Err(e) => {
let _ = cleanup_dir(&out_dir).await;
return Err(CompileError::Failed {
log: format!("latexmk exited successfully but the PDF was not found ({e})"),
});
}
};
// The .fls should always exist under -recorder; degrade gracefully to
// an empty string if missing — the deps-hash will then fall back to
// the path-key, which still caches correctly for self-contained docs.
let fls_text = tokio::fs::read_to_string(&fls_path).await.unwrap_or_default();
let _ = cleanup_dir(&out_dir).await;
Ok((pdf_bytes, fls_text))
}
}
impl Default for LatexCompiler {
fn default() -> Self { Self::new() }
}
// ── Helpers ─────────────────────────────────────────────────────────────────
/// First 5 bytes (10 hex chars) of SHA-256 — enough to avoid collisions in
/// practice while keeping cache filenames short.
fn content_hash(bytes: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(bytes);
let digest = hasher.finalize();
digest.iter().take(5).map(|b| format!("{b:02x}")).collect()
}
/// Short hash of the source file's absolute path. Used to find the `.fls`
/// sidecar that records the dependency list for that source. The path itself
/// (not its content) is hashed so the sidecar location is stable across
/// content edits.
fn path_hash(tex_path: &Path) -> String {
// Canonicalise when possible so that `./foo.tex` and `/abs/foo.tex` resolve
// to the same key. If the file does not exist yet we fall back to the raw
// bytes — path_hash is only ever called for files we are about to compile,
// so this branch is essentially unreachable in practice.
let key: Vec<u8> = std::fs::canonicalize(tex_path)
.map(|p| p.to_string_lossy().into_owned().into_bytes())
.unwrap_or_else(|_| tex_path.to_string_lossy().into_owned().into_bytes());
content_hash(&key)
}
/// Hash of every input file's contents, combined deterministically. Order is
/// stabilised by sorting the paths before hashing so reordering lines in the
/// `.fls` does not invalidate the cache.
///
/// Returns `Err` if any dependency cannot be read — callers should treat that
/// as a cache miss (a `\input` was probably deleted).
async fn composite_hash_of(deps: &[PathBuf]) -> std::io::Result<String> {
let mut sorted: Vec<&PathBuf> = deps.iter().collect();
sorted.sort();
let mut hasher = Sha256::new();
for dep in &sorted {
let bytes = tokio::fs::read(dep).await?;
// Include the path in the hash too: two swapped files with identical
// contents (e.g. chapter1.tex ↔ chapter2.tex) would otherwise collide.
hasher.update(dep.to_string_lossy().as_bytes());
hasher.update(b"\0");
hasher.update(&bytes);
hasher.update(b"\0");
}
let digest = hasher.finalize();
Ok(digest.iter().take(5).map(|b| format!("{b:02x}")).collect())
}
/// Parse a `.fls` recorder file and return the user-controlled input files.
///
/// The `.fls` format is a sequence of `INPUT <path>` and `OUTPUT <path>` lines
/// produced by TeX's `-recorder` flag (orchestrated here by `latexmk`). We
/// keep only the `INPUT` lines that:
///
/// - Are not part of the system TeX distribution (see [`TEXMF_PREFIXES`]);
/// - Are not generated artefacts (see [`AUX_EXTS`]);
/// - Do not live inside the scratch output directory produced during compile.
///
/// `tex_path` provides the CWD that `latexmk` was invoked from, so that
/// relative paths in the `.fls` (always relative to the CWD, not the source
/// file) can be resolved.
fn parse_user_deps(fls_text: &str, tex_path: &Path) -> Vec<PathBuf> {
let cwd = tex_path.parent().unwrap_or_else(|| Path::new("."));
let mut deps: Vec<PathBuf> = Vec::new();
for line in fls_text.lines() {
let path_str = match line.strip_prefix("INPUT ") {
Some(p) => p.trim(),
None => continue,
};
if path_str.is_empty() { continue; }
// Resolve relative paths against the CWD that latexmk was invoked from.
let raw = Path::new(path_str);
let resolved: PathBuf = if raw.is_absolute() {
raw.to_path_buf()
} else {
cwd.join(raw)
};
if !is_user_input(&resolved) { continue; }
if !deps.contains(&resolved) {
deps.push(resolved);
}
}
deps
}
/// Decide whether a file recorded as an INPUT in the `.fls` is a
/// user-controlled dependency worth hashing. See [`TEXMF_PREFIXES`] and
/// [`AUX_EXTS`].
fn is_user_input(path: &Path) -> bool {
let s = path.to_string_lossy();
// Skip anything inside a known TeX distribution prefix.
if TEXMF_PREFIXES.iter().any(|prefix| s.starts_with(prefix)) {
return false;
}
// Skip aux/output artefacts by extension. Handle compound extensions like
// `synctex.gz` by checking the last two segments joined by '.'.
let single_ext = path.extension().and_then(|e| e.to_str());
let compound_ext = path.file_name().and_then(|n| n.to_str()).and_then(|name| {
let parts: Vec<&str> = name.split('.').collect();
if parts.len() >= 2 {
Some(format!("{}.{}", parts[parts.len() - 2], parts[parts.len() - 1]))
} else {
None
}
});
let candidates: [Option<&str>; 2] = [single_ext, compound_ext.as_deref()];
for ext in candidates.into_iter().flatten() {
if AUX_EXTS.iter().any(|aux| *aux == ext) {
return false;
}
}
true
}
/// Per-compile unique suffix (PID + nanosecond timestamp) to namespace the
/// scratch output directory and avoid races between concurrent compiles of the
/// same source.
fn unique_suffix() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let pid = std::process::id();
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
format!("{pid}-{nanos:x}")
}
/// Return the absolute path of `bin` if it is found on `PATH` and is a regular
/// file. We avoid pulling in the `which` crate for a single lookup.
async fn find_on_path(bin: &str) -> Option<PathBuf> {
let path_var = std::env::var_os("PATH")?;
for dir in std::env::split_paths(&path_var) {
let candidate = dir.join(bin);
if tokio::fs::metadata(&candidate).await
.map(|m| m.is_file() || m.file_type().is_symlink())
.unwrap_or(false)
{
return Some(candidate);
}
}
None
}
/// Read `latexmk`'s `.log` from the scratch directory, falling back to a
/// synthetic message if the log is missing or unreadable.
async fn read_compile_log(out_dir: &Path, tex_path: &Path) -> String {
let stem = file_stem(tex_path).unwrap_or_else(|| "output".to_string());
let log_path = out_dir.join(format!("{stem}.log"));
tokio::fs::read_to_string(&log_path)
.await
.unwrap_or_else(|_| String::from("(no log file available)"))
}
/// Recursively remove a scratch directory. Errors are logged and swallowed:
/// leftover dirs only consume a little disk under the OS temp folder.
async fn cleanup_dir(dir: &Path) -> std::io::Result<()> {
if tokio::fs::try_exists(dir).await.unwrap_or(false) {
tokio::fs::remove_dir_all(dir).await?;
}
Ok(())
}
fn file_stem(path: &Path) -> Option<String> {
path.file_stem()
.and_then(|s| s.to_str())
.map(|s| s.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn hash_is_10_lowercase_hex_chars() {
let h = content_hash(b"hello world");
assert_eq!(h.len(), 10);
assert!(h.chars().all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase()));
}
#[test]
fn hash_is_deterministic() {
assert_eq!(content_hash(b"abc"), content_hash(b"abc"));
assert_ne!(content_hash(b"abc"), content_hash(b"abd"));
}
#[test]
fn path_hash_is_stable_for_same_path() {
let p = Path::new("/tmp/foo.tex");
assert_eq!(path_hash(p), path_hash(p));
}
#[test]
fn path_hash_differs_for_different_paths() {
let a = path_hash(Path::new("/tmp/foo.tex"));
let b = path_hash(Path::new("/tmp/bar.tex"));
assert_ne!(a, b);
}
const SAMPLE_FLS: &str = "\
INPUT /usr/local/texlive/2024/texmf-dist/tex/latex/base/article.cls
INPUT chapters/intro.tex
INPUT chapters/intro.tex
INPUT images/diagram.pdf
INPUT refs.bib
INPUT custom.sty
INPUT /Library/TeX/texmf/tex/latex/amsmath/amsmath.sty
INPUT hello.aux
INPUT hello.fls
INPUT hello.synctex.gz
OUTPUT hello.pdf
OUTPUT hello.log
";
#[test]
fn parse_user_deps_filters_texmf_and_aux() {
let deps = parse_user_deps(SAMPLE_FLS, Path::new("/project/hello.tex"));
let mut got: Vec<String> = deps.into_iter()
.map(|p| p.to_string_lossy().into_owned())
.collect();
got.sort();
// Only user-controlled inputs remain, deduped; relative paths resolve
// against the .tex's parent directory.
let mut expected = vec![
"/project/chapters/intro.tex".to_string(),
"/project/images/diagram.pdf".to_string(),
"/project/refs.bib".to_string(),
"/project/custom.sty".to_string(),
];
expected.sort();
assert_eq!(got, expected);
}
#[test]
fn is_user_input_rejects_known_aux_extensions() {
for ext in ["aux", "log", "fls", "fdb_latexmk", "toc", "bbl", "synctex.gz"] {
let path_str = format!("/tmp/out/file.{ext}");
let path = Path::new(&path_str);
assert!(!is_user_input(path), "expected {ext} to be filtered out");
}
}
#[test]
fn is_user_input_keeps_user_files() {
for ext in ["tex", "sty", "cls", "bib", "png", "jpg", "pdf", "eps"] {
let path_str = format!("/project/file.{ext}");
let path = Path::new(&path_str);
assert!(is_user_input(path), "expected {ext} to be kept");
}
}
#[test]
fn is_user_input_rejects_texmf_paths() {
for prefix in TEXMF_PREFIXES {
let path_str = format!("{prefix}/2024/texmf-dist/foo.sty");
let path = Path::new(&path_str);
assert!(!is_user_input(path), "expected {prefix} to be filtered");
}
}
/// Two real files to hash. These used to be `file!()` and `Cargo.toml`,
/// which quietly stopped resolving when this module moved into its own
/// crate: `file!()` is relative to the workspace root while a test's working
/// directory is the package root. Owning the fixtures keeps the test about
/// hashing.
fn hash_fixture(tag: &str) -> (PathBuf, PathBuf, PathBuf) {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
let dir = std::env::temp_dir().join(format!("skald-hash-{tag}-{}-{nanos}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let (a, b) = (dir.join("a.tex"), dir.join("b.tex"));
std::fs::write(&a, b"\\documentclass{article}").unwrap();
std::fs::write(&b, b"\\usepackage{amsmath}").unwrap();
(dir, a, b)
}
#[tokio::test]
async fn composite_hash_is_deterministic_for_same_contents() {
let (dir, a, _b) = hash_fixture("deterministic");
let deps = vec![a];
assert_eq!(
composite_hash_of(&deps).await.unwrap(),
composite_hash_of(&deps).await.unwrap()
);
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn composite_hash_is_order_independent() {
let (dir, a, b) = hash_fixture("order");
let forwards = vec![a.clone(), b.clone()];
let backwards = vec![b, a];
assert_eq!(
composite_hash_of(&forwards).await.unwrap(),
composite_hash_of(&backwards).await.unwrap()
);
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test]
async fn composite_hash_fails_when_a_dep_is_missing() {
let deps = vec![PathBuf::from("/this/path/does/not/exist.tex")];
assert!(composite_hash_of(&deps).await.is_err());
}
#[test]
fn watch_paths_for_returns_just_tex_when_no_fls_cached() {
// Point the compiler at an empty cache directory so no .fls exists.
let compiler = LatexCompiler { cache_dir: PathBuf::from("/tmp/skald-latex-test-empty") };
let tex = Path::new("/project/hello.tex");
let paths = compiler.watch_paths_for(tex);
assert_eq!(paths, vec![tex]);
}
#[test]
fn watch_paths_for_includes_deps_when_fls_is_present() {
// Write a .fls sidecar at the cache location watch_paths_for consults.
// The sidecar's path is derived from path_hash(tex), which we compute
// via the public surface by re-using the same helper.
let tmp = std::env::temp_dir().join(format!("skald-latex-test-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&tmp);
std::fs::create_dir_all(&tmp).unwrap();
let tex = std::env::current_dir().unwrap().join("Cargo.toml"); // arbitrary existing file
let tex_canonical = std::fs::canonicalize(&tex).unwrap();
let path_key = {
let mut h = Sha256::new();
h.update(tex_canonical.to_string_lossy().as_bytes());
let d = h.finalize();
d.iter().take(5).map(|b| format!("{b:02x}")).collect::<String>()
};
let fls_path = tmp.join(format!("{path_key}.fls"));
std::fs::write(&fls_path, format!(
"INPUT /usr/local/texlive/2024/texmf-dist/tex/latex/base/article.cls\n\
INPUT chapters/intro.tex\n\
INPUT custom.sty\n\
INPUT hello.aux\n\
OUTPUT hello.pdf\n\
")
).unwrap();
let compiler = LatexCompiler { cache_dir: tmp.clone() };
let paths = compiler.watch_paths_for(&tex_canonical);
// tex itself + the two user-controlled deps (intro.tex, custom.sty).
// The texmf path and the .aux are filtered out.
assert!(paths.contains(&tex_canonical));
let intro = tex_canonical.parent().unwrap().join("chapters/intro.tex");
let sty = tex_canonical.parent().unwrap().join("custom.sty");
assert!(paths.contains(&intro), "missing {intro:?} in {paths:?}");
assert!(paths.contains(&sty), "missing {sty:?} in {paths:?}");
assert_eq!(paths.len(), 3);
let _ = std::fs::remove_dir_all(&tmp);
}
}
+10
View File
@@ -0,0 +1,10 @@
//! LaTeX → PDF compilation service.
//!
//! Used by the file viewer (`GET /api/file?…&compile-latex=true`) to render
//! `.tex` sources into PDFs on demand. Compilation is delegated to `latexmk`
//! (xelatex engine); results are cached on disk keyed by a short SHA-256 of the
//! source content, so unchanged files are served without recompiling.
pub mod compiler;
pub use compiler::{CompileError, LatexCompiler};
+47
View File
@@ -0,0 +1,47 @@
//! The headless Skald core: storage, identity, LLM stack, tools, sessions.
//!
//! Nothing here knows what runs it. The process shell — HTTP server, desktop
//! webview, setup wizard — lives in the crates that depend on this one. Concrete
//! plugins are never named: `plugin::PluginManager` only ever sees
//! `Arc<dyn Plugin>`, constructed by the consumer and handed to `Skald::new`.
pub mod boot;
pub mod config;
pub mod config_store;
pub mod skald;
pub mod agents;
pub mod approval;
pub mod chat_event_bus;
pub mod chat_hub;
pub mod chatbot;
pub mod clarification;
pub mod command;
pub mod compactor;
pub mod crypto;
pub mod elicitation;
pub mod cron;
pub mod db;
pub mod events;
pub mod image_generate;
pub mod inbox;
pub mod latex;
pub mod llm;
pub mod location;
pub mod memory;
pub mod mcp;
pub mod notification;
pub mod pending_registry;
pub mod plugin;
pub mod projects;
pub mod provider;
pub mod run_context;
pub mod secrets;
pub mod service_manager;
pub mod session;
pub mod tic;
pub mod tool_catalog;
pub mod tool_discovery;
pub mod tools;
pub mod transcribe;
pub mod tts;
pub mod users;
+316
View File
@@ -0,0 +1,316 @@
use anyhow::{Context, Result};
use sqlx::SqlitePool;
use core_api::provider::LlmStrength;
use super::{LlmModelRecord, LlmProviderRecord};
// ── Provider rows ─────────────────────────────────────────────────────────────
#[derive(sqlx::FromRow)]
struct ProviderRow {
id: i64,
name: String,
r#type: String,
api_key: Option<String>,
base_url: Option<String>,
description: Option<String>,
}
pub async fn load_all_providers(pool: &SqlitePool) -> Result<Vec<LlmProviderRecord>> {
let rows = sqlx::query_as::<_, ProviderRow>(
"SELECT id, name, type, api_key, base_url, description FROM llm_providers WHERE removed_at IS NULL ORDER BY name ASC",
)
.fetch_all(pool)
.await
.context("llm_providers: load_all")?;
rows.into_iter().map(provider_row_to_record).collect()
}
pub async fn insert_provider(pool: &SqlitePool, r: &LlmProviderRecord) -> Result<i64> {
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO llm_providers (name, type, api_key, base_url, description)
VALUES (?1, ?2, ?3, ?4, ?5)
RETURNING id",
)
.bind(&r.name)
.bind(&r.provider)
.bind(&r.api_key)
.bind(&r.base_url)
.bind(&r.description)
.fetch_one(pool)
.await
.context("llm_providers: insert")?;
Ok(id)
}
pub async fn update_provider(pool: &SqlitePool, id: i64, r: &LlmProviderRecord) -> Result<()> {
sqlx::query(
"UPDATE llm_providers
SET name=?1, type=?2, api_key=?3, base_url=?4, description=?5
WHERE id=?6",
)
.bind(&r.name)
.bind(&r.provider)
.bind(&r.api_key)
.bind(&r.base_url)
.bind(&r.description)
.bind(id)
.execute(pool)
.await
.context("llm_providers: update")?;
Ok(())
}
pub async fn delete_provider(pool: &SqlitePool, id: i64) -> Result<()> {
// Cascade soft-delete all models belonging to this provider.
sqlx::query(
"UPDATE llm_models SET removed_at = datetime('now') WHERE provider_id = ?1 AND removed_at IS NULL",
)
.bind(id)
.execute(pool)
.await
.context("llm_models: cascade soft-delete for provider")?;
// Remove the API key and mark the provider removed.
sqlx::query(
"UPDATE llm_providers SET removed_at = datetime('now'), api_key = NULL WHERE id = ?1",
)
.bind(id)
.execute(pool)
.await
.context("llm_providers: soft-delete")?;
Ok(())
}
// ── Model rows ────────────────────────────────────────────────────────────────
#[derive(sqlx::FromRow)]
struct ModelRow {
id: i64,
provider_id: i64,
model_id: String,
name: String,
strength: Option<String>,
scope: String,
is_default: i64,
priority: i64,
extra_params: Option<String>,
context_length: Option<i64>,
max_output_tokens: Option<i64>,
knowledge_cutoff: Option<String>,
capabilities: String,
reasoning: Option<String>,
}
pub async fn load_all_models(pool: &SqlitePool) -> Result<Vec<LlmModelRecord>> {
let rows = sqlx::query_as::<_, ModelRow>(
"SELECT id, provider_id, model_id, name, strength, scope, is_default, priority, extra_params,
context_length, max_output_tokens, knowledge_cutoff, capabilities, reasoning
FROM llm_models
WHERE removed_at IS NULL
ORDER BY priority ASC, name ASC",
)
.fetch_all(pool)
.await
.context("llm_models: load_all")?;
rows.into_iter().map(model_row_to_record).collect()
}
pub async fn insert_model(pool: &SqlitePool, r: &LlmModelRecord) -> Result<i64> {
let scope = serde_json::to_string(&r.scope)?;
let extra_params = r.extra_params.as_ref().map(|v| v.to_string());
let capabilities = serde_json::to_string(&r.capabilities)?;
let reasoning = r.reasoning.as_ref().map(|v| v.to_string());
// A model row is never hard-deleted, only soft-deleted via `removed_at`:
// history and telemetry name it, and `llm_requests` keeps only the model's
// name, so the row is what maps that name back to a provider. The unique
// identity is `name` (also the resolution key), so re-adding a previously
// removed model with the same alias would collide with the lingering
// soft-deleted row. Upsert on `name` so that existing row is revived
// (removed_at cleared) and every field overwritten.
let id = sqlx::query_scalar::<_, i64>(
"INSERT INTO llm_models (provider_id, model_id, name, strength, scope, is_default, priority, extra_params,
context_length, max_output_tokens, knowledge_cutoff, capabilities, reasoning)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)
ON CONFLICT(name) DO UPDATE SET
provider_id = excluded.provider_id,
model_id = excluded.model_id,
strength = excluded.strength,
scope = excluded.scope,
is_default = excluded.is_default,
priority = excluded.priority,
extra_params = excluded.extra_params,
context_length = excluded.context_length,
max_output_tokens = excluded.max_output_tokens,
knowledge_cutoff = excluded.knowledge_cutoff,
capabilities = excluded.capabilities,
reasoning = excluded.reasoning,
removed_at = NULL
RETURNING id",
)
.bind(r.provider_id)
.bind(&r.model_id)
.bind(&r.name)
.bind(r.strength.map(strength_str))
.bind(scope)
.bind(r.is_default as i64)
.bind(r.priority as i64)
.bind(extra_params)
.bind(r.context_length)
.bind(r.max_output_tokens)
.bind(&r.knowledge_cutoff)
.bind(capabilities)
.bind(reasoning)
.fetch_one(pool)
.await
.context("llm_models: insert")?;
Ok(id)
}
pub async fn update_model(pool: &SqlitePool, id: i64, r: &LlmModelRecord) -> Result<()> {
let scope = serde_json::to_string(&r.scope)?;
let extra_params = r.extra_params.as_ref().map(|v| v.to_string());
let capabilities = serde_json::to_string(&r.capabilities)?;
let reasoning = r.reasoning.as_ref().map(|v| v.to_string());
sqlx::query(
"UPDATE llm_models
SET provider_id=?1, model_id=?2, name=?3, strength=?4,
scope=?5, is_default=?6, priority=?7, extra_params=?8,
context_length=?9, max_output_tokens=?10, knowledge_cutoff=?11, capabilities=?12,
reasoning=?13
WHERE id=?14",
)
.bind(r.provider_id)
.bind(&r.model_id)
.bind(&r.name)
.bind(r.strength.map(strength_str))
.bind(scope)
.bind(r.is_default as i64)
.bind(r.priority as i64)
.bind(extra_params)
.bind(r.context_length)
.bind(r.max_output_tokens)
.bind(&r.knowledge_cutoff)
.bind(capabilities)
.bind(reasoning)
.bind(id)
.execute(pool)
.await
.context("llm_models: update")?;
Ok(())
}
pub async fn delete_model(pool: &SqlitePool, id: i64) -> Result<()> {
sqlx::query("UPDATE llm_models SET removed_at = datetime('now') WHERE id = ?1")
.bind(id)
.execute(pool)
.await
.context("llm_models: soft-delete")?;
Ok(())
}
/// Update catalog-sourced metadata for a model identified by `provider_id` and `model_id`.
/// Used by the sync logic in `LlmManager::list_provider_models`.
pub async fn update_model_metadata(
pool: &SqlitePool,
provider_id: i64,
model_id: &str,
context_length: Option<i64>,
max_output_tokens: Option<i64>,
knowledge_cutoff: Option<&str>,
capabilities: &[String],
) -> Result<()> {
let caps = serde_json::to_string(capabilities)?;
sqlx::query(
"UPDATE llm_models
SET context_length = COALESCE(?1, context_length),
max_output_tokens = COALESCE(?2, max_output_tokens),
knowledge_cutoff = COALESCE(?3, knowledge_cutoff),
capabilities = ?4
WHERE provider_id = ?5 AND model_id = ?6 AND removed_at IS NULL",
)
.bind(context_length)
.bind(max_output_tokens)
.bind(knowledge_cutoff)
.bind(caps)
.bind(provider_id)
.bind(model_id)
.execute(pool)
.await
.context("llm_models: update_model_metadata")?;
Ok(())
}
pub async fn clear_default(pool: &SqlitePool) -> Result<()> {
sqlx::query("UPDATE llm_models SET is_default=0")
.execute(pool)
.await
.context("llm_models: clear_default")?;
Ok(())
}
// ── Helpers ───────────────────────────────────────────────────────────────────
fn provider_row_to_record(r: ProviderRow) -> Result<LlmProviderRecord> {
Ok(LlmProviderRecord {
id: r.id,
name: r.name,
provider: r.r#type,
api_key: r.api_key,
base_url: r.base_url,
description: r.description,
})
}
fn model_row_to_record(r: ModelRow) -> Result<LlmModelRecord> {
let scope: Vec<String> = serde_json::from_str(&r.scope).unwrap_or_default();
let extra_params = r.extra_params
.as_deref()
.and_then(|s| serde_json::from_str(s).ok());
let capabilities: Vec<String> = serde_json::from_str(&r.capabilities).unwrap_or_default();
let reasoning = r.reasoning
.as_deref()
.and_then(|s| serde_json::from_str(s).ok());
Ok(LlmModelRecord {
id: r.id,
provider_id: r.provider_id,
model_id: r.model_id,
name: r.name,
strength: r.strength.as_deref().and_then(parse_strength),
scope,
is_default: r.is_default != 0,
priority: r.priority as i32,
extra_params,
context_length: r.context_length,
max_output_tokens: r.max_output_tokens,
knowledge_cutoff: r.knowledge_cutoff,
capabilities,
reasoning,
})
}
pub fn strength_str(s: LlmStrength) -> &'static str {
match s {
LlmStrength::VeryLow => "very_low",
LlmStrength::Low => "low",
LlmStrength::Average => "average",
LlmStrength::High => "high",
LlmStrength::VeryHigh => "very_high",
}
}
fn parse_strength(s: &str) -> Option<LlmStrength> {
match s {
"very_low" => Some(LlmStrength::VeryLow),
"low" => Some(LlmStrength::Low),
"average" => Some(LlmStrength::Average),
"high" => Some(LlmStrength::High),
"very_high" => Some(LlmStrength::VeryHigh),
_ => None,
}
}
+582
View File
@@ -0,0 +1,582 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use anyhow::{Context, Result};
use indexmap::IndexMap;
use sqlx::SqlitePool;
use tokio::sync::RwLock;
use tracing::{info, warn};
use crate::chatbot::ChatbotClient;
use crate::chatbot::logging::{LoggingChatbotClient, LogSaveFlags};
use core_api::provider::LlmStrength;
use crate::provider::{ApiProvider, ProviderRegistry, ReasoningMode};
use super::providers::RemoteLlmModelInfo;
use super::{ClientStatus, LlmEntry, LlmModelInfo, LlmModelRecord, LlmProviderInfo, LlmProviderRecord};
use super::db;
const FAILURE_DEGRADED: u32 = 3;
const FAILURE_DOWN: u32 = 5;
const CATALOG_TTL: Duration = Duration::from_secs(24 * 60 * 60);
const MODEL_META_TTL: Duration = Duration::from_secs(60 * 60); // 1 hour
pub const AUTO_CLIENT: &str = "auto";
struct CachedCatalog {
models: Vec<RemoteLlmModelInfo>,
fetched_at: Instant,
}
struct CachedModelMeta {
info: RemoteLlmModelInfo,
fetched_at: Instant,
}
struct HealthState {
status: ClientStatus,
consecutive_failures: u32,
last_error: Option<String>,
}
impl Default for HealthState {
fn default() -> Self {
Self { status: ClientStatus::Healthy, consecutive_failures: 0, last_error: None }
}
}
struct ModelSlot {
provider: LlmProviderRecord,
model: LlmModelRecord,
entry: Arc<LlmEntry>,
health: HealthState,
}
struct ManagerState {
/// Keyed by model.name, ordered by priority ASC.
models: IndexMap<String, ModelSlot>,
/// Keyed by provider.id.
providers: IndexMap<i64, LlmProviderRecord>,
default: String,
}
pub struct LlmManager {
pool: Arc<SqlitePool>,
registry: Arc<ProviderRegistry>,
state: RwLock<ManagerState>,
/// In-memory model catalog cache, keyed by provider_id. TTL = 24h.
catalog: RwLock<HashMap<i64, CachedCatalog>>,
/// Per-model metadata cache, keyed by model display name. TTL = 1h.
model_meta_cache: RwLock<HashMap<String, CachedModelMeta>>,
/// When `Some`, every LLM entry is wrapped with [`LoggingChatbotClient`].
log_flags: Option<LogSaveFlags>,
}
impl LlmManager {
pub async fn new(
pool: Arc<SqlitePool>,
registry: Arc<ProviderRegistry>,
log_flags: Option<LogSaveFlags>,
) -> Result<Arc<Self>> {
let mgr = Arc::new(Self {
pool,
registry,
state: RwLock::new(ManagerState {
models: IndexMap::new(),
providers: IndexMap::new(),
default: String::new(),
}),
catalog: RwLock::new(HashMap::new()),
model_meta_cache: RwLock::new(HashMap::new()),
log_flags,
});
mgr.reload().await?;
Ok(mgr)
}
// ── Public: resolution ────────────────────────────────────────────────────
pub async fn resolve(
&self,
client_name: Option<&str>,
required_scope: Option<&str>,
required_strength: Option<LlmStrength>,
) -> Result<(String, Arc<LlmEntry>)> {
let name = match client_name {
None | Some(AUTO_CLIENT) => {
let (name, entry) = self.select(required_scope, required_strength).await?;
self.maybe_refresh_meta(&name).await;
return Ok((name, entry));
}
Some(n) => {
let state = self.state.read().await;
if !state.models.contains_key(n) {
anyhow::bail!("LLM model '{n}' not found");
}
n.to_string()
}
};
self.maybe_refresh_meta(&name).await;
let state = self.state.read().await;
let entry = state.models.get(&name).map(|s| s.entry.clone())
.with_context(|| format!("LLM model '{name}' not found after refresh"))?;
Ok((name, entry))
}
/// If the per-model metadata cache is stale (or missing) for `name`,
/// fetch fresh data from the provider and update the entry if successful.
async fn maybe_refresh_meta(&self, name: &str) {
{
let cache = self.model_meta_cache.read().await;
if let Some(entry) = cache.get(name) {
if entry.fetched_at.elapsed() < MODEL_META_TTL {
return;
}
}
}
let (provider_id, model_id) = {
let state = self.state.read().await;
match state.models.get(name) {
Some(slot) => (slot.provider.id, slot.model.model_id.clone()),
None => return,
}
};
let remote: RemoteLlmModelInfo = match self.fetch_model_info(provider_id, &model_id).await {
Some(m) => m,
None => return,
};
let now = Instant::now();
let mut cache = self.model_meta_cache.write().await;
cache.insert(name.to_string(), CachedModelMeta { info: remote.clone(), fetched_at: now });
if let Some(ctx) = remote.context_length {
let mut state = self.state.write().await;
if let Some(slot) = state.models.get_mut(name) {
let old_ctx = slot.entry.context_length;
if Some(ctx as i64) != old_ctx {
slot.entry = Arc::new(LlmEntry {
context_length: Some(ctx as i64),
..(*slot.entry).clone()
});
}
}
}
}
async fn fetch_model_info(&self, provider_id: i64, model_id: &str) -> Option<RemoteLlmModelInfo> {
let record = self.state.read().await.providers.get(&provider_id).cloned()?;
let provider = self.registry.get(&record.provider)?;
provider.llm_model_info(&record, model_id).await.ok().flatten()
}
pub async fn get(&self, name: &str) -> Option<Arc<LlmEntry>> {
self.state.read().await.models.get(name).map(|s| s.entry.clone())
}
pub async fn default_name(&self) -> String {
self.state.read().await.default.clone()
}
/// Returns ["auto", <model1>, <model2>, …] for the frontend selector.
pub async fn client_names(&self) -> Vec<String> {
let mut names = vec![AUTO_CLIENT.to_string()];
names.extend(self.state.read().await.models.keys().cloned());
names
}
// ── Public: health reporting ──────────────────────────────────────────────
pub async fn mark_success(&self, name: &str) {
let mut state = self.state.write().await;
if let Some(slot) = state.models.get_mut(name) {
let h = &mut slot.health;
if h.consecutive_failures > 0 {
info!(model = name, "LLM model recovered");
}
h.consecutive_failures = 0;
h.last_error = None;
h.status = ClientStatus::Healthy;
}
}
pub async fn mark_failure(&self, name: &str, error: &str) {
let mut state = self.state.write().await;
if let Some(slot) = state.models.get_mut(name) {
let h = &mut slot.health;
h.consecutive_failures += 1;
h.last_error = Some(error.to_string());
h.status = if h.consecutive_failures >= FAILURE_DOWN {
warn!(model = name, failures = h.consecutive_failures, "LLM model marked DOWN");
ClientStatus::Down
} else if h.consecutive_failures >= FAILURE_DEGRADED {
warn!(model = name, failures = h.consecutive_failures, "LLM model marked DEGRADED");
ClientStatus::Degraded
} else {
ClientStatus::Healthy
};
}
}
// ── Public: provider CRUD ─────────────────────────────────────────────────
pub async fn add_provider(&self, record: LlmProviderRecord) -> Result<i64> {
let id = db::insert_provider(&self.pool, &record).await?;
self.reload().await?;
Ok(id)
}
pub async fn update_provider(&self, id: i64, record: LlmProviderRecord) -> Result<()> {
db::update_provider(&self.pool, id, &record).await?;
self.reload().await
}
pub async fn delete_provider(&self, id: i64) -> Result<()> {
db::delete_provider(&self.pool, id).await?;
self.reload().await
}
pub async fn get_provider(&self, id: i64) -> Option<LlmProviderRecord> {
self.state.read().await.providers.get(&id).cloned()
}
/// Returns the ApiProvider implementation for the given provider record id.
pub async fn get_api_provider(&self, id: i64) -> Option<Arc<dyn ApiProvider>> {
let record = self.state.read().await.providers.get(&id).cloned()?;
self.registry.get(&record.provider)
}
/// Returns the remote model catalog for a provider, using a 24h in-memory cache.
/// After fetching, syncs context/token/capability metadata to existing DB model records.
pub async fn list_provider_models(&self, id: i64) -> Result<Vec<RemoteLlmModelInfo>> {
{
let cache = self.catalog.read().await;
if let Some(entry) = cache.get(&id) {
if entry.fetched_at.elapsed() < CATALOG_TTL {
return Ok(entry.models.clone());
}
}
}
let record = self.state.read().await.providers.get(&id).cloned()
.ok_or_else(|| anyhow::anyhow!("provider {id} not found"))?;
let provider = self.registry.get(&record.provider)
.ok_or_else(|| anyhow::anyhow!("unknown provider type '{}' for provider {id}", record.provider))?;
let mut models = provider.list_llm_models(&record).await?
.ok_or_else(|| anyhow::anyhow!("this provider does not support model listing"))?;
// Fill the reasoning descriptor per catalog model for the add-from-catalog
// UI. Providers that already populate a precise descriptor in their
// listing (e.g. OpenRouter from each model's `reasoning` object) keep it;
// the rest fall back to the capability-based `reasoning_mode`.
for m in &mut models {
if m.reasoning.is_none() {
m.reasoning = provider.reasoning_mode(&m.id, &m.capabilities);
}
}
for remote in &models {
db::update_model_metadata(
&self.pool, id, &remote.id,
remote.context_length.map(|v| v as i64),
remote.max_completion_tokens.map(|v| v as i64),
remote.knowledge_cutoff.as_deref(),
&remote.capabilities,
).await.ok();
}
self.catalog.write().await.insert(id, CachedCatalog {
models: models.clone(),
fetched_at: Instant::now(),
});
Ok(models)
}
pub async fn list_providers_info(&self) -> Vec<LlmProviderInfo> {
self.state.read().await.providers.values().map(|p| {
let supported_types = self.registry.get(&p.provider)
.map(|prov| prov.supported_types().to_vec())
.unwrap_or_default();
LlmProviderInfo {
id: p.id,
name: p.name.clone(),
provider: p.provider.clone(),
base_url: p.base_url.clone(),
description: p.description.clone(),
supported_types,
}
}).collect()
}
// ── Public: model CRUD ────────────────────────────────────────────────────
pub async fn add_model(&self, model: LlmModelRecord) -> Result<i64> {
if model.is_default {
db::clear_default(&self.pool).await?;
}
let id = db::insert_model(&self.pool, &model).await?;
self.reload().await?;
Ok(id)
}
pub async fn update_model(&self, id: i64, model: LlmModelRecord) -> Result<()> {
if model.is_default {
db::clear_default(&self.pool).await?;
}
db::update_model(&self.pool, id, &model).await?;
self.reload().await
}
pub async fn delete_model(&self, id: i64) -> Result<()> {
db::delete_model(&self.pool, id).await?;
self.reload().await
}
pub async fn get_model(&self, id: i64) -> Option<LlmModelRecord> {
self.state.read().await.models.values()
.find(|s| s.model.id == id)
.map(|s| s.model.clone())
}
/// Reasoning control descriptor for a (provider, model_id) pair — used by the
/// "add model" form to render the right control before the model is saved.
/// Capabilities are unknown at this point (manual entry), so an empty slice
/// is passed; catalog-based flows use the descriptor attached to each model.
pub async fn reasoning_mode_for(&self, provider_id: i64, model_id: &str) -> Option<ReasoningMode> {
let record = self.state.read().await.providers.get(&provider_id).cloned()?;
let provider = self.registry.get(&record.provider)?;
provider.reasoning_mode(model_id, &[])
}
pub async fn list_models_info(&self) -> Vec<LlmModelInfo> {
let state = self.state.read().await;
let catalog = self.catalog.read().await;
state.models.values().map(|slot| {
let cached = catalog.get(&slot.provider.id)
.and_then(|c| c.models.iter().find(|m| m.id == slot.model.model_id));
let reasoning_mode = self.registry.get(&slot.provider.provider)
.and_then(|p| p.reasoning_mode(&slot.model.model_id, &slot.model.capabilities));
LlmModelInfo {
id: slot.model.id,
provider_id: slot.provider.id,
provider_name: slot.provider.name.clone(),
model_id: slot.model.model_id.clone(),
name: slot.model.name.clone(),
strength: slot.model.strength,
scope: slot.model.scope.clone(),
is_default: slot.model.is_default,
priority: slot.model.priority,
extra_params: slot.model.extra_params.clone(),
context_length: slot.model.context_length,
max_output_tokens: slot.model.max_output_tokens,
knowledge_cutoff: slot.model.knowledge_cutoff.clone(),
capabilities: slot.model.capabilities.clone(),
status: slot.health.status,
last_error: slot.health.last_error.clone(),
price_input_per_million: cached.and_then(|m| m.price_input_per_million),
price_output_per_million: cached.and_then(|m| m.price_output_per_million),
reasoning: slot.model.reasoning.clone(),
reasoning_mode,
}
}).collect()
}
// ── Public: selection ─────────────────────────────────────────────────────
pub async fn select_excluding(
&self,
excluded: &[&str],
required_scope: Option<&str>,
required_strength: Option<LlmStrength>,
) -> Result<(String, Arc<LlmEntry>)> {
let state = self.state.read().await;
let mut slots: Vec<(&String, &ModelSlot)> = state.models.iter()
.filter(|(name, _)| !excluded.contains(&name.as_str()))
.collect();
if slots.is_empty() {
anyhow::bail!("no alternative LLM models available");
}
sort_slots_for_agent(&mut slots, required_scope, required_strength);
if let Some((name, slot)) = slots.iter().find(|(_, s)| s.health.status != ClientStatus::Down) {
return Ok((name.to_string(), slot.entry.clone()));
}
if let Some((name, slot)) = slots.first() {
warn!(model = %name, "all alternative LLM models are DOWN — using best available");
return Ok((name.to_string(), slot.entry.clone()));
}
anyhow::bail!("no alternative LLM models available");
}
async fn select(
&self,
required_scope: Option<&str>,
required_strength: Option<LlmStrength>,
) -> Result<(String, Arc<LlmEntry>)> {
let state = self.state.read().await;
if state.models.is_empty() {
anyhow::bail!("no LLM models configured — add one via the UI");
}
let mut slots: Vec<(&String, &ModelSlot)> = state.models.iter().collect();
sort_slots_for_agent(&mut slots, required_scope, required_strength);
if let Some((name, slot)) = slots.iter().find(|(_, s)| s.health.status != ClientStatus::Down) {
return Ok((name.to_string(), slot.entry.clone()));
}
if let Some((name, slot)) = slots.first() {
warn!(model = %name, "all LLM models are DOWN — using strongest as emergency fallback");
return Ok((name.to_string(), slot.entry.clone()));
}
anyhow::bail!("no LLM models available");
}
// ── Private ───────────────────────────────────────────────────────────────
async fn reload(&self) -> Result<()> {
let provider_records = db::load_all_providers(&self.pool).await?;
let model_records = db::load_all_models(&self.pool).await?;
let providers: IndexMap<i64, LlmProviderRecord> = provider_records
.into_iter()
.map(|p| (p.id, p))
.collect();
let mut models: IndexMap<String, ModelSlot> = IndexMap::new();
let mut default = String::new();
for model in model_records {
let provider = match providers.get(&model.provider_id) {
Some(p) => p.clone(),
None => {
warn!(model = %model.name, provider_id = model.provider_id, "orphaned model — provider not found, skipping");
continue;
}
};
let log_config = self.log_flags.map(|f| (Arc::clone(&self.pool), f));
let entry = match build_entry(&self.registry, &provider, &model, model.id, log_config) {
Ok(e) => Arc::new(e),
Err(e) => {
warn!(model = %model.name, error = %e, "failed to build LLM entry, skipping");
continue;
}
};
if model.is_default || default.is_empty() {
default = model.name.clone();
}
models.insert(model.name.clone(), ModelSlot {
provider,
model,
entry,
health: HealthState::default(),
});
}
let mut state = self.state.write().await;
for (name, slot) in state.models.iter() {
if let Some(new_slot) = models.get_mut(name) {
new_slot.health.status = slot.health.status;
new_slot.health.consecutive_failures = slot.health.consecutive_failures;
new_slot.health.last_error = slot.health.last_error.clone();
}
}
state.models = models;
state.providers = providers;
state.default = default;
Ok(())
}
}
// ── Builder ───────────────────────────────────────────────────────────────────
fn build_entry(
registry: &ProviderRegistry,
provider: &LlmProviderRecord,
model: &LlmModelRecord,
model_db_id: i64,
log_config: Option<(Arc<SqlitePool>, LogSaveFlags)>,
) -> Result<LlmEntry> {
let built = registry.get(&provider.provider)
.ok_or_else(|| anyhow::anyhow!("unknown provider type '{}'", provider.provider))?
.build_llm(provider, model)
.ok_or_else(|| anyhow::anyhow!("provider '{}' does not support LLM", provider.provider))??;
let inner = built.client;
let prompt_cache = built.prompt_cache;
let extra = model.extra_params.clone();
let client: Arc<dyn ChatbotClient> = match log_config {
Some((pool, flags)) => Arc::new(LoggingChatbotClient::new(inner, pool, &model.name, flags)),
None => inner,
};
Ok(LlmEntry {
client,
model: model.model_id.clone(),
model_db_id,
strength: model.strength,
scope: model.scope.clone(),
extra_params: extra,
context_length: model.context_length,
prompt_cache,
})
}
// ── Sorting helpers ───────────────────────────────────────────────────────────
pub fn sort_models_for_agent(
mut models: Vec<LlmModelInfo>,
scope: Option<&str>,
strength: Option<LlmStrength>,
) -> Vec<LlmModelInfo> {
models.sort_by_key(|m| (model_tier(m.strength, m.scope.as_slice(), scope, strength), m.priority));
models
}
fn sort_slots_for_agent(
slots: &mut Vec<(&String, &ModelSlot)>,
scope: Option<&str>,
strength: Option<LlmStrength>,
) {
slots.sort_by_key(|(_, s)| (
model_tier(s.model.strength, s.model.scope.as_slice(), scope, strength),
s.model.priority,
));
}
fn model_tier(
model_strength: Option<LlmStrength>,
model_scope: &[String],
req_scope: Option<&str>,
req_strength: Option<LlmStrength>,
) -> u8 {
let strength_ok = match (req_strength, model_strength) {
(Some(req), Some(avail)) => avail >= req,
(Some(_), None) => false,
(None, _) => true,
};
// Prefer exact strength match over over-qualified models so that e.g. an
// agent with strength=low picks the `low` model before `average`.
let exact_match = match (req_strength, model_strength) {
(Some(req), Some(avail)) => avail == req,
_ => true,
};
let scope_ok = req_scope.map_or(true, |sc| model_scope.iter().any(|x| x == sc));
match (strength_ok && scope_ok, exact_match && scope_ok, strength_ok) {
(true, true, _) => 0, // exact strength + scope ok
(true, false, _) => 1, // over-qualified but scope ok
(false, _, true) => 2, // strength ok, scope mismatch
_ => 3, // doesn't meet minimum bar
}
}
+82
View File
@@ -0,0 +1,82 @@
pub(crate) mod db;
pub mod manager;
pub mod providers;
use std::sync::Arc;
use crate::chatbot::ChatbotClient;
use crate::provider::ServiceType;
pub use core_api::provider::{LlmProviderRecord, LlmModelRecord, LlmStrength, ReasoningMode};
pub use manager::{LlmManager, sort_models_for_agent};
/// A resolved, ready-to-use LLM client with its associated metadata.
#[derive(Clone)]
pub struct LlmEntry {
pub client: Arc<dyn ChatbotClient>,
pub model: String,
pub model_db_id: i64,
pub strength: Option<LlmStrength>,
pub scope: Vec<String>,
pub extra_params: Option<serde_json::Value>,
/// Max input context window in tokens, if known.
pub context_length: Option<i64>,
/// When true, prompt-caching hints are injected into requests.
pub prompt_cache: bool,
}
// ── Provider ──────────────────────────────────────────────────────────────────
/// Public provider metadata (no api_key).
#[derive(Debug, Clone, serde::Serialize)]
pub struct LlmProviderInfo {
pub id: i64,
pub name: String,
#[serde(rename = "type")]
pub provider: String,
pub base_url: Option<String>,
pub description: Option<String>,
/// Service types this provider supports (from ProviderRegistry at runtime).
pub supported_types: Vec<ServiceType>,
}
/// Public model metadata for API responses (includes provider name for convenience).
#[derive(Debug, Clone, serde::Serialize)]
pub struct LlmModelInfo {
pub id: i64,
pub provider_id: i64,
pub provider_name: String,
pub model_id: String,
pub name: String,
pub strength: Option<LlmStrength>,
pub scope: Vec<String>,
pub is_default: bool,
pub priority: i32,
pub extra_params: Option<serde_json::Value>,
pub context_length: Option<i64>,
pub max_output_tokens: Option<i64>,
pub knowledge_cutoff: Option<String>,
pub capabilities: Vec<String>,
pub status: ClientStatus,
pub last_error: Option<String>,
/// Input (prompt) price per million tokens (USD) from the provider catalog cache.
pub price_input_per_million: Option<f64>,
/// Output (completion) price per million tokens (USD) from the provider catalog cache.
pub price_output_per_million: Option<f64>,
/// Currently-selected reasoning value (string for a `ValueSet`, number for a
/// `Range`, or `None`). Round-trips to the edit form.
pub reasoning: Option<serde_json::Value>,
/// Reasoning control descriptor for this model (drives the UI control), or
/// `None` if the model does not support reasoning.
pub reasoning_mode: Option<ReasoningMode>,
}
// ── Health ────────────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ClientStatus {
Healthy,
Degraded,
Down,
}
@@ -0,0 +1,119 @@
use std::sync::Arc;
use anyhow::{Context, Result, anyhow};
use crate::chatbot::anthropic::AnthropicClient;
use crate::llm::{LlmModelRecord, LlmProviderRecord};
use crate::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
use crate::provider::{ApiProvider, BuiltLlmClient, ProviderField, ProviderUiMeta, ReasoningMode, ServiceType};
pub struct AnthropicProvider {
http: reqwest::Client,
}
impl AnthropicProvider {
pub fn new() -> Self {
Self { http: reqwest::Client::new() }
}
}
#[async_trait::async_trait]
impl ApiProvider for AnthropicProvider {
fn type_id(&self) -> &'static str { "anthropic" }
fn display_name(&self) -> &'static str { "Anthropic" }
fn supported_types(&self) -> &'static [ServiceType] {
&[ServiceType::Llm]
}
async fn list_llm_models(&self, _record: &LlmProviderRecord) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
Ok(None)
}
async fn llm_model_info(&self, record: &LlmProviderRecord, model_id: &str) -> Result<Option<RemoteLlmModelInfo>> {
let api_key = record.api_key.as_deref()
.ok_or_else(|| anyhow!("provider '{}': api_key required for anthropic model_info", record.name))?;
let url = format!("https://api.anthropic.com/v1/models/{model_id}");
let resp: serde_json::Value = self.http
.get(&url)
.header("x-api-key", api_key)
.header("anthropic-version", "2023-06-01")
.send()
.await
.map_err(|e| anyhow!("Anthropic model_info request failed: {e}"))?
.json()
.await
.map_err(|e| anyhow!("Anthropic model_info response parse failed: {e}"))?;
let id = resp["id"].as_str().ok_or_else(|| anyhow!("missing 'id' in Anthropic response"))?.to_string();
let name = resp["display_name"].as_str().unwrap_or(&id).to_string();
Ok(Some(RemoteLlmModelInfo {
id,
name,
context_length: resp["context_window"].as_u64(),
max_completion_tokens: resp["max_output_tokens"].as_u64(),
knowledge_cutoff: None,
capabilities: vec![],
vision: None,
price_input_per_million: None,
price_output_per_million: None,
reasoning: None,
}))
}
fn reasoning_mode(&self, model_id: &str, capabilities: &[String]) -> Option<ReasoningMode> {
// Extended thinking → numeric token budget. Available on Claude 3.7 and
// the 4.x/5.x families (not the 3.5/3-opus generation).
let id = model_id.to_lowercase();
let supports = capabilities.iter().any(|c| c == "reasoning")
|| id.contains("3-7")
|| id.contains("-4") || id.contains("-5")
|| id.contains("opus-4") || id.contains("sonnet-4") || id.contains("haiku-4");
if supports {
Some(ReasoningMode::Range {
min: 1024,
max: 32_000,
step: Some(1024),
default: Some(8192),
unit: Some("tokens".to_string()),
})
} else {
None
}
}
fn reasoning_request(&self, value: &serde_json::Value) -> Option<serde_json::Value> {
// value is a JSON number (budget_tokens).
let budget = value.as_i64().filter(|n| *n > 0)?;
Some(serde_json::json!({
"thinking": { "type": "enabled", "budget_tokens": budget }
}))
}
fn build_llm(&self, record: &LlmProviderRecord, model: &LlmModelRecord) -> Option<Result<BuiltLlmClient>> {
Some((|| {
let key = record.api_key.as_deref()
.with_context(|| format!("provider '{}': api_key required for anthropic", record.name))?;
// Merge model extra_params + reasoning (thinking) into the request body.
let extra = extra_with_reasoning(self, model);
Ok(BuiltLlmClient {
client: Arc::new(AnthropicClient::with_extra_body(key, extra)),
prompt_cache: false,
})
})())
}
fn ui_meta(&self) -> ProviderUiMeta {
ProviderUiMeta {
type_id: "anthropic",
display_name: "Anthropic",
description: None,
color: "#d4a574",
icon: "bi-chat-square-dots",
fields: &[
ProviderField { key: "api_key", label: "API Key", required: true, secret: true },
],
}
}
}
@@ -0,0 +1,146 @@
use std::sync::Arc;
use anyhow::{Context, Result, anyhow};
use crate::chatbot::openai::OpenAiClient;
use crate::llm::{LlmModelRecord, LlmProviderRecord};
use crate::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
use crate::provider::{ApiProvider, BuiltLlmClient, ProviderField, ProviderUiMeta, ReasoningMode, ServiceType};
pub struct DeepSeekProvider {
http: reqwest::Client,
}
impl DeepSeekProvider {
pub fn new() -> Self {
Self { http: reqwest::Client::new() }
}
fn known_context_length(model_id: &str) -> Option<u64> {
let id = model_id.to_lowercase();
if id.contains("coder") { Some(16384) }
else if id.contains("reasoner") { Some(65536) }
else if id.starts_with("deepseek-v4") { Some(1_048_576) }
else if id.starts_with("deepseek-chat") || id.starts_with("deepseek-v3") { Some(65536) }
else { None }
}
fn known_max_output(model_id: &str) -> Option<u64> {
if model_id.to_lowercase().starts_with("deepseek-v4") { Some(393_216) } else { None }
}
fn known_capabilities(model_id: &str) -> Vec<String> {
let mut caps = vec!["function_calling".to_string()];
if model_id.to_lowercase().contains("reasoner") {
caps.push("reasoning".to_string());
}
caps
}
}
#[async_trait::async_trait]
impl ApiProvider for DeepSeekProvider {
fn type_id(&self) -> &'static str { "deepseek" }
fn display_name(&self) -> &'static str { "DeepSeek" }
fn supported_types(&self) -> &'static [ServiceType] {
&[ServiceType::Llm]
}
async fn list_llm_models(&self, record: &LlmProviderRecord) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
let api_key = record.api_key.as_deref()
.ok_or_else(|| anyhow!("provider '{}': api_key required for deepseek model listing", record.name))?;
let resp: serde_json::Value = self.http
.get("https://api.deepseek.com/models")
.bearer_auth(api_key)
.send()
.await
.map_err(|e| anyhow!("DeepSeek request failed: {e}"))?
.error_for_status()
.map_err(|e| anyhow!("DeepSeek error response: {e}"))?
.json()
.await
.map_err(|e| anyhow!("DeepSeek response parse failed: {e}"))?;
let models = resp["data"]
.as_array()
.ok_or_else(|| anyhow!("unexpected DeepSeek response shape"))?
.iter()
.filter_map(|m| {
let id = m["id"].as_str()?.to_string();
let name = id.clone();
let context_length = Self::known_context_length(&id).or_else(|| m["context_length"].as_u64());
let capabilities = Self::known_capabilities(&id);
let max_output = Self::known_max_output(&id);
Some(RemoteLlmModelInfo {
id, name, context_length,
max_completion_tokens: max_output,
knowledge_cutoff: None,
capabilities,
vision: None,
price_input_per_million: None,
price_output_per_million: None,
reasoning: None,
})
})
.collect();
Ok(Some(models))
}
fn reasoning_mode(&self, model_id: &str, capabilities: &[String]) -> Option<ReasoningMode> {
// Thinking mode (thinking.type) + graded reasoning_effort. "disabled"
// turns thinking off; effort levels low/medium map to high, xhigh to max.
let id = model_id.to_lowercase();
if capabilities.iter().any(|c| c == "reasoning")
|| id.contains("reasoner")
|| id.starts_with("deepseek-v4")
{
Some(ReasoningMode::ValueSet {
values: ["disabled", "low", "medium", "high", "xhigh", "max"]
.iter().map(|s| s.to_string()).collect(),
default: Some("high".to_string()),
})
} else {
None
}
}
fn reasoning_request(&self, value: &serde_json::Value) -> Option<serde_json::Value> {
// "disabled" → thinking off; "enabled" → thinking on (no effort);
// any effort level → thinking on + `reasoning_effort`.
match value.as_str()? {
"disabled" => Some(serde_json::json!({ "thinking": { "type": "disabled" } })),
"enabled" => Some(serde_json::json!({ "thinking": { "type": "enabled" } })),
effort => Some(serde_json::json!({
"thinking": { "type": "enabled" },
"reasoning_effort": effort,
})),
}
}
fn build_llm(&self, record: &LlmProviderRecord, model: &LlmModelRecord) -> Option<Result<BuiltLlmClient>> {
Some((|| {
let key = record.api_key.as_deref()
.with_context(|| format!("provider '{}': api_key required for deepseek", record.name))?;
let extra = extra_with_reasoning(self, model);
Ok(BuiltLlmClient {
client: Arc::new(OpenAiClient::new("https://api.deepseek.com/v1", key, extra, false)),
prompt_cache: false,
})
})())
}
fn ui_meta(&self) -> ProviderUiMeta {
ProviderUiMeta {
type_id: "deepseek",
display_name: "DeepSeek",
description: None,
color: "#0ea5e9",
icon: "bi-search",
fields: &[
ProviderField { key: "api_key", label: "API Key", required: true, secret: true },
],
}
}
}
@@ -0,0 +1,86 @@
use std::sync::Arc;
use anyhow::{Result, anyhow};
use crate::chatbot::lm_studio::LmStudioClient;
use crate::llm::{LlmModelRecord, LlmProviderRecord};
use crate::llm::providers::RemoteLlmModelInfo;
use crate::provider::{ApiProvider, BuiltLlmClient, ProviderField, ProviderUiMeta, ServiceType};
pub struct LmStudioProvider {
http: reqwest::Client,
}
impl LmStudioProvider {
pub fn new() -> Self {
Self { http: reqwest::Client::new() }
}
fn base_url(record: &LlmProviderRecord) -> String {
record.base_url.clone()
.unwrap_or_else(|| "http://localhost:1234/v1".to_string())
}
}
#[async_trait::async_trait]
impl ApiProvider for LmStudioProvider {
fn type_id(&self) -> &'static str { "lm_studio" }
fn display_name(&self) -> &'static str { "LM Studio" }
fn supported_types(&self) -> &'static [ServiceType] {
&[ServiceType::Llm]
}
async fn list_llm_models(&self, record: &LlmProviderRecord) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
let url = format!("{}/models", Self::base_url(record).trim_end_matches('/'));
let resp: serde_json::Value = self.http
.get(&url)
.send()
.await
.map_err(|e| anyhow!("LM Studio request failed: {e}"))?
.json()
.await
.map_err(|e| anyhow!("LM Studio response parse failed: {e}"))?;
let models = resp["data"]
.as_array()
.ok_or_else(|| anyhow!("unexpected LM Studio response shape"))?
.iter()
.filter_map(|m| {
let id = m["id"].as_str()?.to_string();
Some(RemoteLlmModelInfo {
name: id.clone(), id,
context_length: None,
max_completion_tokens: None,
knowledge_cutoff: None,
capabilities: vec![],
vision: None,
price_input_per_million: None,
price_output_per_million: None,
reasoning: None,
})
})
.collect();
Ok(Some(models))
}
fn build_llm(&self, record: &LlmProviderRecord, _model: &LlmModelRecord) -> Option<Result<BuiltLlmClient>> {
Some(Ok(BuiltLlmClient {
client: Arc::new(LmStudioClient::new(record.base_url.as_deref())),
prompt_cache: false,
}))
}
fn ui_meta(&self) -> ProviderUiMeta {
ProviderUiMeta {
type_id: "lm_studio",
display_name: "LM Studio",
description: Some("Local models via LM Studio"),
color: "#6b7280",
icon: "bi-window-stack",
fields: &[
ProviderField { key: "base_url", label: "Base URL", required: false, secret: false },
],
}
}
}
@@ -0,0 +1,39 @@
pub mod anthropic;
pub mod deepseek;
pub mod lm_studio;
pub mod ollama;
pub mod openai;
pub mod openrouter;
pub mod zai;
// Re-export so existing code that uses `providers::ServiceType` / `providers::RemoteLlmModelInfo` keeps working.
pub use crate::provider::ServiceType;
pub use core_api::provider::RemoteLlmModelInfo;
use core_api::provider::{ApiProvider, LlmModelRecord};
/// Computes the `extra_params` an OpenAI-compatible client should be built with,
/// given a model's stored `extra_params` and its selected reasoning value. The
/// provider translates the reasoning value into a request fragment via
/// `reasoning_request`; that fragment's top-level keys are merged over
/// `extra_params` (reasoning wins on conflict). Returns `None` when neither is set.
pub(crate) fn extra_with_reasoning(
provider: &dyn ApiProvider,
model: &LlmModelRecord,
) -> Option<serde_json::Value> {
let reasoning = model.reasoning.as_ref().and_then(|v| provider.reasoning_request(v));
match (model.extra_params.clone(), reasoning) {
(base, None) => base,
(None, overlay) => overlay,
(Some(mut base), Some(overlay)) => {
match (base.as_object_mut(), overlay.as_object()) {
(Some(b), Some(o)) => {
for (k, v) in o { b.insert(k.clone(), v.clone()); }
Some(base)
}
// Non-object base: the reasoning overlay takes precedence.
_ => Some(overlay),
}
}
}
}
@@ -0,0 +1,123 @@
use std::sync::Arc;
use anyhow::{Result, anyhow};
use crate::chatbot::ollama::OllamaClient;
use crate::llm::{LlmModelRecord, LlmProviderRecord};
use crate::llm::providers::RemoteLlmModelInfo;
use crate::provider::{ApiProvider, BuiltLlmClient, ProviderField, ProviderUiMeta, ServiceType};
pub struct OllamaProvider {
http: reqwest::Client,
}
impl OllamaProvider {
pub fn new() -> Self {
Self { http: reqwest::Client::new() }
}
fn base_url(record: &LlmProviderRecord) -> String {
record.base_url.clone()
.unwrap_or_else(|| "http://localhost:11434".to_string())
}
fn parse_model_info(show: &serde_json::Value, model_id: &str) -> RemoteLlmModelInfo {
let context_length = show["model_info"]["llm.context_length"]
.as_u64()
.or_else(|| {
show["model_info"]["llm.context_length"]
.as_str()
.and_then(|s| s.parse::<u64>().ok())
});
RemoteLlmModelInfo {
name: model_id.to_string(),
id: model_id.to_string(),
context_length,
max_completion_tokens: None,
knowledge_cutoff: None,
capabilities: vec![],
vision: None,
price_input_per_million: None,
price_output_per_million: None,
reasoning: None,
}
}
}
#[async_trait::async_trait]
impl ApiProvider for OllamaProvider {
fn type_id(&self) -> &'static str { "ollama" }
fn display_name(&self) -> &'static str { "Ollama" }
fn supported_types(&self) -> &'static [ServiceType] {
&[ServiceType::Llm]
}
async fn list_llm_models(&self, record: &LlmProviderRecord) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
let url = format!("{}/api/tags", Self::base_url(record).trim_end_matches('/'));
let resp: serde_json::Value = self.http
.get(&url)
.send()
.await
.map_err(|e| anyhow!("Ollama request failed: {e}"))?
.json()
.await
.map_err(|e| anyhow!("Ollama response parse failed: {e}"))?;
let models = resp["models"]
.as_array()
.ok_or_else(|| anyhow!("unexpected Ollama response shape"))?
.iter()
.filter_map(|m| {
let id = m["name"].as_str()?.to_string();
Some(RemoteLlmModelInfo {
name: id.clone(), id,
context_length: None,
max_completion_tokens: None,
knowledge_cutoff: None,
capabilities: vec![],
vision: None,
price_input_per_million: None,
price_output_per_million: None,
reasoning: None,
})
})
.collect();
Ok(Some(models))
}
async fn llm_model_info(&self, record: &LlmProviderRecord, model_id: &str) -> Result<Option<RemoteLlmModelInfo>> {
let url = format!("{}/api/show", Self::base_url(record).trim_end_matches('/'));
let body = serde_json::json!({ "name": model_id });
let resp: serde_json::Value = self.http
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| anyhow!("Ollama model_info request failed: {e}"))?
.json()
.await
.map_err(|e| anyhow!("Ollama model_info response parse failed: {e}"))?;
Ok(Some(Self::parse_model_info(&resp, model_id)))
}
fn build_llm(&self, record: &LlmProviderRecord, _model: &LlmModelRecord) -> Option<Result<BuiltLlmClient>> {
Some(Ok(BuiltLlmClient {
client: Arc::new(OllamaClient::new(record.base_url.as_deref())),
prompt_cache: false,
}))
}
fn ui_meta(&self) -> ProviderUiMeta {
ProviderUiMeta {
type_id: "ollama",
display_name: "Ollama",
description: Some("Local models via Ollama"),
color: "#f97316",
icon: "bi-terminal",
fields: &[
ProviderField { key: "base_url", label: "Base URL", required: false, secret: false },
],
}
}
}
@@ -0,0 +1,99 @@
use std::sync::Arc;
use anyhow::{Context, Result};
use crate::chatbot::openai::OpenAiClient;
use crate::llm::{LlmModelRecord, LlmProviderRecord};
use crate::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
use crate::transcribe::TranscribeModelRecord;
use crate::transcribe::openai_audio::OpenAiAudioTranscriber;
use crate::tts::TtsModelRecord;
use crate::tts::openai_tts::OpenAiTtsSynthesiser;
use crate::provider::{ApiProvider, BuiltLlmClient, ProviderField, ProviderUiMeta, ReasoningMode, ServiceType};
pub struct OpenAiProvider;
#[async_trait::async_trait]
impl ApiProvider for OpenAiProvider {
fn type_id(&self) -> &'static str { "open_ai" }
fn display_name(&self) -> &'static str { "OpenAI" }
fn supported_types(&self) -> &'static [ServiceType] {
&[ServiceType::Llm, ServiceType::Transcribe, ServiceType::Tts]
}
async fn list_llm_models(&self, _record: &LlmProviderRecord) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
Ok(None)
}
fn reasoning_mode(&self, model_id: &str, capabilities: &[String]) -> Option<ReasoningMode> {
// Reasoning ("o" series and other reasoning models) → effort levels.
let id = model_id.to_lowercase();
let is_reasoning = capabilities.iter().any(|c| c == "reasoning")
|| id.starts_with("o1") || id.starts_with("o3") || id.starts_with("o4")
|| id.starts_with("gpt-5");
if is_reasoning {
Some(ReasoningMode::ValueSet {
values: vec!["low".to_string(), "medium".to_string(), "high".to_string()],
default: Some("medium".to_string()),
})
} else {
None
}
}
fn reasoning_request(&self, value: &serde_json::Value) -> Option<serde_json::Value> {
let effort = value.as_str()?;
Some(serde_json::json!({ "reasoning_effort": effort }))
}
fn build_llm(&self, record: &LlmProviderRecord, model: &LlmModelRecord) -> Option<Result<BuiltLlmClient>> {
Some((|| {
let key = record.api_key.as_deref()
.with_context(|| format!("provider '{}': api_key required for open_ai", record.name))?;
let extra = extra_with_reasoning(self, model);
Ok(BuiltLlmClient {
client: Arc::new(OpenAiClient::new("https://api.openai.com/v1", key, extra, false)),
prompt_cache: false,
})
})())
}
fn build_tts(&self, record: &LlmProviderRecord, model: &TtsModelRecord) -> Option<Result<Arc<dyn crate::tts::TextToSpeech>>> {
Some((|| {
let base_url = record.base_url.clone()
.unwrap_or_else(|| "https://api.openai.com/v1".to_string());
let api_key = record.api_key.clone()
.with_context(|| format!("provider '{}': api_key required for open_ai", record.name))?;
Ok(Arc::new(OpenAiTtsSynthesiser::new(
&model.name, base_url, api_key, &model.model_id,
model.voice_id.clone(), model.instructions.clone(), model.response_format.clone(),
)) as Arc<dyn crate::tts::TextToSpeech>)
})())
}
fn build_transcriber(&self, record: &LlmProviderRecord, model: &TranscribeModelRecord) -> Option<Result<Arc<dyn crate::transcribe::Transcribe>>> {
Some((|| {
let base_url = record.base_url.clone()
.unwrap_or_else(|| "https://api.openai.com/v1".to_string());
let api_key = record.api_key.clone()
.with_context(|| format!("provider '{}': api_key required for open_ai", record.name))?;
Ok(Arc::new(OpenAiAudioTranscriber::new(
&model.name, base_url, api_key, &model.model_id, model.language.clone(),
)) as Arc<dyn crate::transcribe::Transcribe>)
})())
}
fn ui_meta(&self) -> ProviderUiMeta {
ProviderUiMeta {
type_id: "open_ai",
display_name: "OpenAI",
description: None,
color: "#10a37f",
icon: "bi-lightning-charge",
fields: &[
ProviderField { key: "api_key", label: "API Key", required: true, secret: true },
ProviderField { key: "base_url", label: "Base URL (optional)", required: false, secret: false },
],
}
}
}
@@ -0,0 +1,215 @@
use std::sync::Arc;
use anyhow::{Context, Result, anyhow};
use crate::chatbot::openai::OpenAiClient;
use crate::image_generate::ImageGenerateModelRecord;
use crate::image_generate::openrouter_image::OpenRouterImageGenerator;
use crate::llm::{LlmModelRecord, LlmProviderRecord};
use crate::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
use crate::transcribe::TranscribeModelRecord;
use crate::transcribe::openai_audio::OpenAiAudioTranscriber;
use crate::tts::TtsModelRecord;
use crate::tts::openai_tts::OpenAiTtsSynthesiser;
use crate::provider::{ApiProvider, BuiltLlmClient, ProviderField, ProviderUiMeta, ReasoningMode, ServiceType};
pub struct OpenRouterProvider {
http: reqwest::Client,
}
impl OpenRouterProvider {
pub fn new() -> Self {
Self { http: reqwest::Client::new() }
}
/// Builds a `ReasoningMode` from OpenRouter's per-model `reasoning` object
/// (`/api/v1/models`): `supported_efforts` → discrete effort levels,
/// otherwise `supports_max_tokens` → a token-budget range. Non-reasoning
/// models omit the object entirely → `None`. The UI's "— off —" option
/// (a null value → no reasoning param sent) covers disabling.
fn parse_reasoning(v: &serde_json::Value) -> Option<ReasoningMode> {
if !v.is_object() {
return None;
}
let default = v["default_effort"].as_str().map(String::from);
let efforts: Vec<String> = v["supported_efforts"].as_array()
.map(|a| a.iter().filter_map(|e| e.as_str().map(String::from)).collect())
.unwrap_or_default();
if !efforts.is_empty() {
Some(ReasoningMode::ValueSet { values: efforts, default })
} else if v["supports_max_tokens"].as_bool().unwrap_or(false) {
Some(ReasoningMode::Range {
min: 1024, max: 32_000, step: Some(1024), default: Some(8192),
unit: Some("tokens".to_string()),
})
} else {
None
}
}
async fn fetch_catalog(&self, api_key: &str) -> Result<Vec<RemoteLlmModelInfo>> {
let resp: serde_json::Value = self.http
.get("https://openrouter.ai/api/v1/models")
.bearer_auth(api_key)
.send()
.await
.map_err(|e| anyhow!("OpenRouter request failed: {e}"))?
.json()
.await
.map_err(|e| anyhow!("OpenRouter response parse failed: {e}"))?;
let models = resp["data"]
.as_array()
.ok_or_else(|| anyhow!("unexpected OpenRouter response shape"))?
.iter()
.filter_map(|m| {
let id = m["id"].as_str()?.to_string();
let name = m["name"].as_str().unwrap_or(&id).to_string();
let context_length = m["context_length"].as_u64();
let price_input = m["pricing"]["prompt"].as_str()
.and_then(|s| s.parse::<f64>().ok())
.map(|v| v * 1_000_000.0);
let price_output = m["pricing"]["completion"].as_str()
.and_then(|s| s.parse::<f64>().ok())
.map(|v| v * 1_000_000.0);
let capabilities = {
let mut caps = vec!["function_calling".to_string()];
if let Some(params) = m["supported_parameters"].as_array() {
for p in params {
if let Some(s) = p.as_str() {
match s {
"tools" => caps.push("function_calling".to_string()),
"vision" | "image" => caps.push("vision".to_string()),
"stream" => caps.push("streaming".to_string()),
"reasoning" | "reasoning_effort" => caps.push("reasoning".to_string()),
_ => {}
}
}
}
}
caps.sort();
caps.dedup();
caps
};
let vision = Some(capabilities.contains(&"vision".to_string()));
let reasoning = Self::parse_reasoning(&m["reasoning"]);
Some(RemoteLlmModelInfo {
id, name, context_length,
max_completion_tokens: None,
knowledge_cutoff: None,
capabilities,
vision,
price_input_per_million: price_input,
price_output_per_million: price_output,
reasoning,
})
})
.collect();
Ok(models)
}
}
#[async_trait::async_trait]
impl ApiProvider for OpenRouterProvider {
fn type_id(&self) -> &'static str { "openrouter" }
fn display_name(&self) -> &'static str { "OpenRouter" }
fn supported_types(&self) -> &'static [ServiceType] {
&[ServiceType::Llm, ServiceType::Transcribe, ServiceType::ImageGenerate, ServiceType::Tts]
}
async fn list_llm_models(&self, record: &LlmProviderRecord) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
let api_key = record.api_key.as_deref()
.ok_or_else(|| anyhow!("provider '{}': api_key required for openrouter model listing", record.name))?;
Ok(Some(self.fetch_catalog(api_key).await?))
}
fn reasoning_mode(&self, _model_id: &str, capabilities: &[String]) -> Option<ReasoningMode> {
// Fallback for stored/manually-added models with no catalog descriptor.
// The precise per-model set comes from `parse_reasoning` in the catalog;
// here we offer OpenRouter's full accepted effort set.
if capabilities.iter().any(|c| c == "reasoning") {
Some(ReasoningMode::ValueSet {
values: ["minimal", "low", "medium", "high", "xhigh", "max"]
.iter().map(|s| s.to_string()).collect(),
default: Some("medium".to_string()),
})
} else {
None
}
}
fn reasoning_request(&self, value: &serde_json::Value) -> Option<serde_json::Value> {
// String → effort level; number → token budget (max_tokens).
if let Some(effort) = value.as_str() {
Some(serde_json::json!({ "reasoning": { "effort": effort } }))
} else {
let budget = value.as_i64().filter(|n| *n > 0)?;
Some(serde_json::json!({ "reasoning": { "max_tokens": budget } }))
}
}
fn build_llm(&self, record: &LlmProviderRecord, model: &LlmModelRecord) -> Option<Result<BuiltLlmClient>> {
Some((|| {
let key = record.api_key.as_deref()
.with_context(|| format!("provider '{}': api_key required for openrouter", record.name))?;
// Anthropic prompt-caching only works for models served by Anthropic on OpenRouter.
let prompt_cache = model.model_id.starts_with("anthropic/");
let extra = extra_with_reasoning(self, model);
Ok(BuiltLlmClient {
client: Arc::new(OpenAiClient::new("https://openrouter.ai/api/v1", key, extra, prompt_cache)),
prompt_cache,
})
})())
}
fn build_tts(&self, record: &LlmProviderRecord, model: &TtsModelRecord) -> Option<Result<Arc<dyn crate::tts::TextToSpeech>>> {
Some((|| {
let base_url = record.base_url.clone()
.unwrap_or_else(|| "https://openrouter.ai/api/v1".to_string());
let api_key = record.api_key.clone()
.with_context(|| format!("provider '{}': api_key required for openrouter", record.name))?;
Ok(Arc::new(OpenAiTtsSynthesiser::new(
&model.name, base_url, api_key, &model.model_id,
model.voice_id.clone(), model.instructions.clone(), model.response_format.clone(),
)) as Arc<dyn crate::tts::TextToSpeech>)
})())
}
fn build_transcriber(&self, record: &LlmProviderRecord, model: &TranscribeModelRecord) -> Option<Result<Arc<dyn crate::transcribe::Transcribe>>> {
Some((|| {
let base_url = record.base_url.clone()
.unwrap_or_else(|| "https://openrouter.ai/api/v1".to_string());
let api_key = record.api_key.clone()
.with_context(|| format!("provider '{}': api_key required for openrouter", record.name))?;
Ok(Arc::new(OpenAiAudioTranscriber::new(
&model.name, base_url, api_key, &model.model_id, model.language.clone(),
)) as Arc<dyn crate::transcribe::Transcribe>)
})())
}
fn build_image_generator(&self, record: &LlmProviderRecord, model: &ImageGenerateModelRecord) -> Option<Result<Arc<dyn crate::image_generate::ImageGenerate>>> {
Some((|| {
let base_url = record.base_url.clone()
.unwrap_or_else(|| "https://openrouter.ai/api/v1".to_string());
let api_key = record.api_key.clone()
.with_context(|| format!("provider '{}': api_key required for openrouter", record.name))?;
Ok(Arc::new(OpenRouterImageGenerator::new(
&model.name, base_url, api_key, &model.model_id,
)) as Arc<dyn crate::image_generate::ImageGenerate>)
})())
}
fn ui_meta(&self) -> ProviderUiMeta {
ProviderUiMeta {
type_id: "openrouter",
display_name: "OpenRouter",
description: None,
color: "#8b5cf6",
icon: "bi-hdd-stack",
fields: &[
ProviderField { key: "api_key", label: "API Key", required: true, secret: true },
],
}
}
}
+145
View File
@@ -0,0 +1,145 @@
use std::sync::Arc;
use anyhow::{Context, Result};
use crate::chatbot::openai::OpenAiClient;
use crate::llm::{LlmModelRecord, LlmProviderRecord};
use crate::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
use crate::provider::{ApiProvider, BuiltLlmClient, ProviderField, ProviderUiMeta, ReasoningMode, ServiceType};
/// Z.AI (Zhipu AI) — OpenAI-compatible GLM API.
///
/// Endpoint `https://api.z.ai/api/paas/v4/chat/completions`; `OpenAiClient`
/// appends `/chat/completions`, so the base URL is `.../paas/v4`.
///
/// Z.AI exposes no `GET /models` endpoint, so the model catalog is a curated
/// static list of the currently published GLM models.
pub struct ZaiProvider;
impl ZaiProvider {
pub fn new() -> Self {
Self
}
/// Base URL for the OpenAI-compatible chat endpoint (without `/chat/completions`).
const BASE_URL: &'static str = "https://api.z.ai/api/paas/v4";
/// Curated GLM catalog. Z.AI has no `GET /models` endpoint; this mirrors the
/// model menu published on the Z.AI console.
fn catalog() -> &'static [&'static str] {
&[
"glm-5.2",
"glm-5.1",
"glm-5",
"glm-5-turbo",
"glm-4.7",
"glm-4.6",
"glm-4.5",
"glm-4-32b-0414-128k",
]
}
fn known_context_length(model_id: &str) -> Option<u64> {
let id = model_id.to_lowercase();
if id.contains("128k") { Some(131_072) }
else if id.starts_with("glm-5") { Some(1_048_576) } // GLM-5.x: 1M context (per Z.AI)
else if id.starts_with("glm-4.7") { Some(200_000) }
else if id.starts_with("glm-4.6") { Some(200_000) }
else if id.starts_with("glm-4.5") { Some(131_072) }
else { None }
}
}
#[async_trait::async_trait]
impl ApiProvider for ZaiProvider {
fn type_id(&self) -> &'static str { "zai" }
fn display_name(&self) -> &'static str { "Z.AI" }
fn supported_types(&self) -> &'static [ServiceType] {
&[ServiceType::Llm]
}
async fn list_llm_models(&self, _record: &LlmProviderRecord) -> Result<Option<Vec<RemoteLlmModelInfo>>> {
let models = Self::catalog()
.iter()
.map(|id| RemoteLlmModelInfo {
id: id.to_string(),
name: id.to_string(),
context_length: Self::known_context_length(id),
max_completion_tokens: None,
knowledge_cutoff: None,
capabilities: vec!["function_calling".to_string()],
vision: Some(false),
price_input_per_million: None,
price_output_per_million: None,
reasoning: None,
})
.collect();
Ok(Some(models))
}
fn reasoning_mode(&self, model_id: &str, _capabilities: &[String]) -> Option<ReasoningMode> {
let id = model_id.to_lowercase();
// GLM-5.2 (and above) additionally expose a graded `reasoning_effort`
// on top of the thinking toggle, so offer the effort levels directly
// ("disabled" turns thinking off).
if id.starts_with("glm-5.2") {
Some(ReasoningMode::ValueSet {
values: ["disabled", "minimal", "low", "medium", "high", "xhigh", "max"]
.iter().map(|s| s.to_string()).collect(),
default: Some("max".to_string()),
})
// Deep-thinking toggle (thinking.type) is supported by the GLM-5.x
// series and GLM-4.5/4.6/4.7 (but not the older glm-4-32b).
} else if id.starts_with("glm-5")
|| id.starts_with("glm-4.7")
|| id.starts_with("glm-4.6")
|| id.starts_with("glm-4.5")
{
Some(ReasoningMode::ValueSet {
values: vec!["disabled".to_string(), "enabled".to_string()],
default: Some("enabled".to_string()),
})
} else {
None
}
}
fn reasoning_request(&self, value: &serde_json::Value) -> Option<serde_json::Value> {
// "disabled" → thinking off; "enabled" → thinking on (no effort);
// any effort level → thinking on + `reasoning_effort` (GLM-5.2+).
match value.as_str()? {
"disabled" => Some(serde_json::json!({ "thinking": { "type": "disabled" } })),
"enabled" => Some(serde_json::json!({ "thinking": { "type": "enabled" } })),
effort => Some(serde_json::json!({
"thinking": { "type": "enabled" },
"reasoning_effort": effort,
})),
}
}
fn build_llm(&self, record: &LlmProviderRecord, model: &LlmModelRecord) -> Option<Result<BuiltLlmClient>> {
Some((|| {
let key = record.api_key.as_deref()
.with_context(|| format!("provider '{}': api_key required for zai", record.name))?;
let extra = extra_with_reasoning(self, model);
Ok(BuiltLlmClient {
client: Arc::new(OpenAiClient::new(Self::BASE_URL, key, extra, false)),
prompt_cache: false,
})
})())
}
fn ui_meta(&self) -> ProviderUiMeta {
ProviderUiMeta {
type_id: "zai",
display_name: "Z.AI",
description: Some("Zhipu AI GLM models (OpenAI-compatible)"),
color: "#4f46e5",
icon: "bi-stars",
fields: &[
ProviderField { key: "api_key", label: "API Key", required: true, secret: true },
],
}
}
}
+1
View File
@@ -0,0 +1 @@
pub use core_api::location::{GpsCoord, LocationEntry, LocationManager, LocationUpdater};
+189
View File
@@ -0,0 +1,189 @@
//! Per-server MCP log files.
//!
//! Consumes [`McpLogLine`]s emitted by the MCP client crate and appends each to a
//! dedicated file `logs/mcp/<name>.log`. Sources captured (see `crates/mcp-client`):
//! child `stderr` (stdio), diverted `notifications/message` log records, and
//! connection lifecycle events. No SQLite — a plain file per server, meant to be
//! scanned later (e.g. by a diagnostics agent) for `[error]`/`[warning]` lines.
use std::collections::HashMap;
use std::path::PathBuf;
use mcp_client::McpLogLine;
use tokio::fs::{File, OpenOptions};
use tokio::io::AsyncWriteExt;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
/// Size at which a server's `.log` is rotated to `.log.1` (one backup kept), so a
/// chatty server can't grow its file without bound.
const MAX_LOG_BYTES: u64 = 5 * 1024 * 1024;
/// Background task: drain `rx` and append every line to its server's file until
/// shutdown or the channel closes. Spawned once from `McpManager::new`.
pub(super) async fn log_consumer(
mut rx: mpsc::UnboundedReceiver<McpLogLine>,
shutdown: CancellationToken,
) {
let mut writer = LogWriter::new(PathBuf::from("logs").join("mcp"));
if let Err(e) = tokio::fs::create_dir_all(&writer.dir).await {
warn!("mcp logs: cannot create {}: {e}", writer.dir.display());
return;
}
info!("mcp: per-server log consumer started ({})", writer.dir.display());
loop {
tokio::select! {
_ = shutdown.cancelled() => {
info!("mcp: log consumer shutdown");
break;
}
msg = rx.recv() => match msg {
Some(line) => writer.write_line(line).await,
None => break,
}
}
}
}
/// Holds one append handle per server (plus its tracked byte size for rotation).
struct LogWriter {
dir: PathBuf,
files: HashMap<String, (File, u64)>,
}
impl LogWriter {
fn new(dir: PathBuf) -> Self {
Self { dir, files: HashMap::new() }
}
/// `logs/mcp/<sanitized>.log`. The name is sanitized so a server called
/// `foo/bar` or `a b` can't escape the directory or produce an odd filename.
fn path_for(&self, server: &str) -> PathBuf {
self.dir.join(format!("{}.log", sanitize(server)))
}
async fn open(&self, server: &str) -> std::io::Result<(File, u64)> {
let path = self.path_for(server);
let file = OpenOptions::new().create(true).append(true).open(&path).await?;
// Seed the tracked size from the existing file so appends keep counting
// toward the rotation threshold across restarts.
let size = file.metadata().await.map(|m| m.len()).unwrap_or(0);
Ok((file, size))
}
/// Renames `<name>.log` → `<name>.log.1` (overwriting any previous backup) and
/// reopens a fresh, empty handle.
async fn rotate(&self, server: &str) -> std::io::Result<(File, u64)> {
let path = self.path_for(server);
let backup = PathBuf::from(format!("{}.1", path.display()));
let _ = tokio::fs::rename(&path, &backup).await; // best-effort
self.open(server).await
}
async fn write_line(&mut self, line: McpLogLine) {
let record = format_record(&line);
let bytes = record.len() as u64;
// Open on first use for this server.
if !self.files.contains_key(&line.server) {
match self.open(&line.server).await {
Ok(handle) => { self.files.insert(line.server.clone(), handle); }
Err(e) => { warn!("mcp logs: open failed for '{}': {e}", line.server); return; }
}
}
// Rotate before writing if this line would push the file over the cap.
let over_cap = self.files.get(&line.server)
.map(|(_, sz)| *sz + bytes > MAX_LOG_BYTES)
.unwrap_or(false);
if over_cap {
match self.rotate(&line.server).await {
Ok(handle) => { self.files.insert(line.server.clone(), handle); }
Err(e) => { warn!("mcp logs: rotate failed for '{}': {e}", line.server); }
}
}
if let Some((file, size)) = self.files.get_mut(&line.server) {
match file.write_all(record.as_bytes()).await {
Ok(()) => {
*size += bytes;
let _ = file.flush().await;
}
Err(e) => {
warn!("mcp logs: write failed for '{}': {e}", line.server);
// Drop the handle so the next line retries a fresh open.
self.files.remove(&line.server);
}
}
}
}
}
/// `2026-07-03T12:34:56.789Z [warning] <text>` — an ISO-8601 UTC timestamp, the
/// padded level tag, then the text. The padded tag keeps files column-aligned and
/// makes `[error]`/`[warning]` trivial to grep.
fn format_record(line: &McpLogLine) -> String {
let ts = chrono::Utc::now().format("%Y-%m-%dT%H:%M:%S%.3fZ");
let tag = format!("[{}]", line.level);
format!("{ts} {tag:<12} {}\n", line.text)
}
/// Keeps ASCII alphanumerics and `-_.`; everything else becomes `_`. Guarantees a
/// non-empty, path-separator-free filename component.
fn sanitize(name: &str) -> String {
let s: String = name.chars()
.map(|c| if c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.') { c } else { '_' })
.collect();
if s.is_empty() { "unknown".to_string() } else { s }
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn sanitize_replaces_separators_and_spaces() {
assert_eq!(sanitize("gmail"), "gmail");
assert_eq!(sanitize("foo/bar"), "foo_bar");
assert_eq!(sanitize("a b"), "a_b");
assert_eq!(sanitize("claude_ai_Gmail"), "claude_ai_Gmail");
assert_eq!(sanitize(""), "unknown");
}
#[test]
fn format_record_has_timestamp_tag_and_text() {
let rec = format_record(&McpLogLine::stderr("srv", "hello"));
assert!(rec.contains("[stderr]"));
assert!(rec.ends_with("hello\n"));
assert!(rec.starts_with("20")); // year prefix of the ISO timestamp
}
#[tokio::test]
async fn write_line_appends_one_file_per_server() {
let dir = std::env::temp_dir().join(format!("skald_mcplogs_{}", std::process::id()));
let _ = tokio::fs::remove_dir_all(&dir).await;
tokio::fs::create_dir_all(&dir).await.unwrap();
let mut writer = LogWriter::new(dir.clone());
writer.write_line(McpLogLine::stderr("gmail", "banner line")).await;
writer.write_line(McpLogLine::lifecycle("gmail", "connected — 3 tool(s)")).await;
writer.write_line(McpLogLine::from_message(
"firecrawl",
&serde_json::json!({ "level": "error", "data": "boom" }),
)).await;
// One file per server, named from the sanitized server name.
let gmail = tokio::fs::read_to_string(dir.join("gmail.log")).await.unwrap();
assert!(gmail.contains("[stderr]") && gmail.contains("banner line"));
assert!(gmail.contains("[lifecycle]") && gmail.contains("connected — 3 tool(s)"));
let fire = tokio::fs::read_to_string(dir.join("firecrawl.log")).await.unwrap();
assert!(fire.contains("[error]") && fire.contains("boom"));
// gmail's lines must not leak into firecrawl's file.
assert!(!fire.contains("banner line"));
let _ = tokio::fs::remove_dir_all(&dir).await;
}
}
+457
View File
@@ -0,0 +1,457 @@
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use anyhow::Result;
use rand::RngExt;
use serde_json::{Value, json};
use sqlx::SqlitePool;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use crate::tools::ToolResult;
pub use mcp_client::{
ElicitationHandler,
McpCallResult, McpLogLine, McpLogTx, McpMedia, McpMediaData, McpMediaKind,
McpServerClient, McpServerConfig, McpServerInfo, McpServerStatus, McpTool, McpTransport as McpTransportKind,
parse_mcp_tool_name,
http_server::McpHttpServer,
server::{McpNotification, McpServer},
};
use mcp_client::McpTransport;
mod logs;
const SERVER_START_TIMEOUT_SECS: u64 = 120;
// ── McpManager ───────────────────────────────────────────────────────────────
pub struct McpManager {
pool: Arc<SqlitePool>,
servers: RwLock<HashMap<String, Arc<dyn McpServerClient>>>,
errors: RwLock<HashMap<String, String>>,
descriptions: RwLock<HashMap<String, Option<String>>>,
notification_tx: mpsc::UnboundedSender<McpNotification>,
/// Feeds per-server diagnostic lines (stderr, `notifications/message`,
/// lifecycle) to the `logs::log_consumer`, which writes `logs/mcp/<name>.log`.
log_tx: McpLogTx,
/// Bridges server-initiated `elicitation/create` requests to the Inbox.
/// Set once via `set_elicitation_handler` before `initialize` runs.
elicitation_handler: RwLock<Option<Arc<dyn ElicitationHandler>>>,
/// Data root for persisting non-text tool-result media (`media_dir`).
data_root: PathBuf,
}
impl McpManager {
pub fn new(pool: Arc<SqlitePool>, shutdown: CancellationToken, data_root: impl Into<PathBuf>) -> Self {
let (notification_tx, notification_rx) = mpsc::unbounded_channel::<McpNotification>();
let (log_tx, log_rx) = mpsc::unbounded_channel::<McpLogLine>();
let pool_bg = pool.clone();
tokio::spawn(Self::notification_consumer(pool_bg, notification_rx, shutdown.clone()));
tokio::spawn(logs::log_consumer(log_rx, shutdown));
Self {
pool,
servers: RwLock::new(HashMap::new()),
errors: RwLock::new(HashMap::new()),
descriptions: RwLock::new(HashMap::new()),
notification_tx,
log_tx,
elicitation_handler: RwLock::new(None),
data_root: data_root.into(),
}
}
/// Emits a lifecycle line to a server's per-server log file (start failure,
/// timeout, connection). Used for transports that have no `stderr` of their
/// own to carry connection diagnostics (notably HTTP/SSE).
fn log_lifecycle(&self, server: &str, text: impl Into<String>) {
let _ = self.log_tx.send(McpLogLine::lifecycle(server.to_string(), text));
}
/// Directory under the data root where inline tool-result media (images,
/// audio, embedded resources) is persisted and served from `/api/mcp-media/`.
pub fn media_dir(&self) -> PathBuf {
self.data_root.join("mcp_media")
}
/// Wire the elicitation bridge. Must be called before `initialize` so that
/// stdio servers are started with a handler for `elicitation/create`.
pub fn set_elicitation_handler(&self, handler: Arc<dyn ElicitationHandler>) {
*self.elicitation_handler.write().unwrap() = Some(handler);
}
fn elicitation_handler(&self) -> Option<Arc<dyn ElicitationHandler>> {
self.elicitation_handler.read().unwrap().clone()
}
async fn notification_consumer(
pool: Arc<SqlitePool>,
mut rx: mpsc::UnboundedReceiver<McpNotification>,
shutdown: CancellationToken,
) {
loop {
tokio::select! {
_ = shutdown.cancelled() => {
info!("mcp: notification consumer shutdown");
break;
}
msg = rx.recv() => match msg {
Some((source, payload)) => {
let method = payload["method"].as_str().unwrap_or("unknown").to_string();
let params = serde_json::to_string(&payload["params"]).unwrap_or_else(|_| "{}".to_string());
match crate::db::mcp_events::insert(&pool, &source, &method, &params).await {
Ok(id) => info!("mcp_event stored: id={id} source={source} method={method}"),
Err(e) => warn!("mcp_events insert failed (source={source} method={method}): {e}"),
}
}
None => break,
}
}
}
}
fn cfg_from_row(row: &crate::db::mcp_servers::McpServerRow) -> McpServerConfig {
McpServerConfig {
name: row.name.clone(),
transport: match row.transport.as_str() {
"http" => McpTransport::Http,
"sse" => McpTransport::Sse,
_ => McpTransport::Stdio,
},
command: row.command.clone(),
args: Some(row.args()).filter(|v| !v.is_empty()),
env: Some(row.env()).filter(|m| !m.is_empty()),
url: row.url.clone(),
api_key: row.api_key.clone(),
}
}
async fn start_one(
cfg: &McpServerConfig,
notification_tx: Option<mpsc::UnboundedSender<McpNotification>>,
log_tx: Option<McpLogTx>,
elicitation_handler: Option<Arc<dyn ElicitationHandler>>,
) -> Result<Arc<dyn McpServerClient>> {
match cfg.transport {
McpTransport::Stdio => {
// Elicitation and per-server diagnostic capture (stderr +
// notifications/message) are stdio-only. HTTP/SSE has no stderr and
// no async notification stream, so it only gets lifecycle lines,
// emitted by the manager (see `log_lifecycle`).
McpServer::start(cfg, notification_tx, log_tx, elicitation_handler).await
.map(|s| Arc::new(s) as Arc<dyn McpServerClient>)
}
McpTransport::Http | McpTransport::Sse => {
McpHttpServer::start(cfg).await
.map(|s| Arc::new(s) as Arc<dyn McpServerClient>)
}
}
}
pub async fn initialize(&self) {
let rows = match crate::db::mcp_servers::all_enabled(&self.pool).await {
Ok(r) => r,
Err(e) => { warn!("McpManager::initialize: failed to read DB: {e}"); return; }
};
if rows.is_empty() {
info!("No enabled MCP servers in DB — MCP disabled.");
crate::boot::section("MCP servers — none enabled");
return;
}
let cfgs: Vec<_> = rows.iter().map(Self::cfg_from_row).collect();
{
let mut descs = self.descriptions.write().unwrap();
for row in &rows {
descs.insert(row.name.clone(), row.description.clone());
}
}
crate::boot::section(format!(
"MCP servers — connecting to {} in background", cfgs.len()
));
let handles: Vec<_> = cfgs.into_iter().map(|cfg| {
let tx = self.notification_tx.clone();
let log_tx = self.log_tx.clone();
let eh = self.elicitation_handler();
tokio::spawn(async move {
info!("MCP server '{}': starting…", cfg.name);
let result = tokio::time::timeout(
Duration::from_secs(SERVER_START_TIMEOUT_SECS),
Self::start_one(&cfg, Some(tx), Some(log_tx), eh),
).await;
(cfg.name, cfg.transport, result)
})
}).collect();
for handle in handles {
match handle.await {
Ok((name, _, Ok(Ok(s)))) => {
let tool_names: Vec<_> = s.tools().iter().map(|t| t.name.as_str()).collect();
info!("MCP server '{}' ready — {} tool(s): {}", name, tool_names.len(), tool_names.join(", "));
let n = tool_names.len();
crate::boot::ok(format!("{name} ({n} tool{})", if n == 1 { "" } else { "s" }));
self.log_lifecycle(&name, format!("connected — {n} tool(s)"));
self.servers.write().unwrap().insert(name, s);
}
Ok((name, _, Ok(Err(e)))) => {
warn!("MCP server '{}' failed to start: {e}", name);
crate::boot::fail(format!("{name}{e}"));
self.log_lifecycle(&name, format!("failed to start: {e}"));
self.errors.write().unwrap().insert(name, e.to_string());
}
Ok((name, _, Err(_))) => {
let msg = format!("startup timed out after {SERVER_START_TIMEOUT_SECS}s");
warn!("MCP server '{}' {msg}", name);
crate::boot::fail(format!("{name}{msg}"));
self.log_lifecycle(&name, &msg);
self.errors.write().unwrap().insert(name, msg);
}
Err(e) => { warn!("MCP startup task panicked: {e}"); }
}
}
}
pub async fn register(&self, p: crate::db::mcp_servers::UpsertParams<'_>) -> Result<Vec<String>> {
let name = p.name.to_string();
crate::db::mcp_servers::upsert(&self.pool, p).await?;
let rows = crate::db::mcp_servers::all_enabled(&self.pool).await?;
let row = rows.into_iter().find(|r| r.name == name)
.ok_or_else(|| anyhow::anyhow!("register: server '{}' not found after upsert", name))?;
let cfg = Self::cfg_from_row(&row);
let client = tokio::time::timeout(
Duration::from_secs(SERVER_START_TIMEOUT_SECS),
Self::start_one(&cfg, Some(self.notification_tx.clone()), Some(self.log_tx.clone()), self.elicitation_handler()),
).await
.map_err(|_| {
self.log_lifecycle(&name, "timed out during connection");
anyhow::anyhow!("MCP server '{}' timed out during connection", name)
})?
.map_err(|e| {
self.log_lifecycle(&name, format!("failed to start: {e}"));
anyhow::anyhow!("MCP server '{}' failed to start: {e}", name)
})?;
let tool_names: Vec<String> = client.tools().iter().map(|t| t.name.clone()).collect();
self.log_lifecycle(&name, format!("connected — {} tool(s)", tool_names.len()));
self.errors.write().unwrap().remove(&name);
self.descriptions.write().unwrap().insert(name.clone(), row.description.clone());
self.servers.write().unwrap().insert(name, client);
Ok(tool_names)
}
pub async fn unregister(&self, name: &str) -> Result<()> {
crate::db::mcp_servers::delete(&self.pool, name).await?;
self.servers.write().unwrap().remove(name);
self.errors.write().unwrap().remove(name);
self.descriptions.write().unwrap().remove(name);
Ok(())
}
pub async fn set_enabled(&self, name: &str, enabled: bool) -> Result<()> {
crate::db::mcp_servers::set_enabled(&self.pool, name, enabled).await
}
pub async fn list(&self) -> Result<Vec<McpServerInfo>> {
let rows = crate::db::mcp_servers::all(&self.pool).await?;
let servers = self.servers.read().unwrap();
let errors = self.errors.read().unwrap();
let infos = rows.into_iter().map(|row| {
let status = if !row.enabled {
McpServerStatus::Disabled
} else if let Some(s) = servers.get(&row.name) {
McpServerStatus::Running {
tools: s.tools().iter().map(|t| t.name.clone()).collect(),
}
} else if let Some(e) = errors.get(&row.name) {
McpServerStatus::Error { message: e.clone() }
} else {
McpServerStatus::Error { message: "not connected".to_string() }
};
McpServerInfo {
name: row.name,
transport: row.transport,
description: row.description,
friendly_name: row.friendly_name,
status,
}
}).collect();
Ok(infos)
}
pub fn tools(&self) -> Vec<McpTool> {
self.servers.read().unwrap().values()
.flat_map(|s| s.tools().iter().cloned())
.collect()
}
pub fn tools_for(&self, names: &[String]) -> Vec<McpTool> {
self.servers.read().unwrap().iter()
.filter(|(name, _)| names.contains(name))
.flat_map(|(_, s)| s.tools().iter().cloned())
.collect()
}
pub fn server_descriptions(&self) -> HashMap<String, Option<String>> {
self.descriptions.read().unwrap().clone()
}
pub fn server_infos(&self) -> Vec<Value> {
self.servers.read().unwrap().iter()
.map(|(name, s)| json!({
"name": name,
"tools": s.tools().iter().map(|t| json!({
"name": t.name,
"description": t.description,
})).collect::<Vec<_>>(),
}))
.collect()
}
pub async fn call(&self, server: &str, tool: &str, args: Value) -> Result<ToolResult> {
let s = self.servers.read().unwrap()
.get(server)
.cloned()
.ok_or_else(|| anyhow::anyhow!("MCP server '{server}' not found"))?;
match s.call_tool(tool, args).await? {
McpCallResult::Text(t) => Ok(ToolResult::Text(t)),
McpCallResult::Json(v) => Ok(ToolResult::Json(v)),
McpCallResult::Media { text, structured, items } =>
Ok(ToolResult::Text(self.persist_media(server, text, structured, items).await)),
// Experimental Tasks — defensive fallback. Normally the transport's
// `call_tool` polls a deferred task to completion (block-and-poll) and
// returns the real result, so this arm is not hit. It only surfaces a
// raw handle if polling was bypassed, so the result is never lost.
McpCallResult::Task(t) => {
let ttl = t.ttl_ms.map(|ms| format!(", ttl {}s", ms / 1000)).unwrap_or_default();
Ok(ToolResult::Text(format!(
"MCP server '{server}' deferred this call as task `{}` (status: {:?}{ttl}). \
Task polling is not implemented yet, so the result can't be retrieved automatically.",
t.task_id, t.status,
)))
}
}
}
/// Persists the inline media of an MCP tool result under [`media_dir`] and
/// composes a markdown text result that references each item by URL — so the
/// model can surface it (the frontend renders the markdown) instead of the
/// bytes being silently dropped. `resource_link`s are passed through by URI
/// without downloading. Falls back to a textual placeholder if a write fails,
/// so a disk error never loses the rest of the result.
async fn persist_media(
&self,
server: &str,
text: Option<String>,
structured: Option<Value>,
items: Vec<McpMedia>,
) -> String {
let mut out: Vec<String> = Vec::new();
if let Some(t) = text.filter(|t| !t.is_empty()) {
out.push(t);
}
for item in items {
match item.data {
McpMediaData::Inline { bytes, mime } => {
let file = format!("{}.{}", random_id(), ext_for_mime(&mime));
let dir = self.media_dir();
let saved = async {
tokio::fs::create_dir_all(&dir).await?;
tokio::fs::write(dir.join(&file), &bytes).await
}.await;
match saved {
Ok(()) => {
let url = format!("/api/mcp-media/{file}");
let kb = bytes.len().div_ceil(1024);
out.push(match item.kind {
McpMediaKind::Image => format!("![image]({url}) ({mime}, {kb} KB)"),
McpMediaKind::Audio => format!("[audio]({url}) ({mime}, {kb} KB)"),
McpMediaKind::Resource => format!("[file]({url}) ({mime}, {kb} KB)"),
});
}
Err(e) => {
warn!("MCP '{server}': failed to persist tool-result media: {e}");
out.push(format!("[media not saved: {mime}]"));
}
}
}
McpMediaData::Link { uri, mime } => {
let label = mime.as_deref().unwrap_or("resource");
out.push(format!("[{label}]({uri})"));
}
}
}
if let Some(sc) = structured {
if let Ok(s) = serde_json::to_string_pretty(&sc) {
out.push(format!("```json\n{s}\n```"));
}
}
out.join("\n\n")
}
}
/// Generates a 32-char alphanumeric id for a persisted media filename
/// (mirrors `ImageGeneratorManager`).
fn random_id() -> String {
rand::rng()
.sample_iter(rand::distr::Alphanumeric)
.take(32)
.map(char::from)
.collect()
}
/// Maps a MIME type to a file extension for persisted MCP media; `bin` for unknown.
pub fn ext_for_mime(mime: &str) -> &'static str {
match mime.split(';').next().unwrap_or("").trim() {
"image/png" => "png",
"image/jpeg" => "jpg",
"image/gif" => "gif",
"image/webp" => "webp",
"image/svg+xml" => "svg",
"audio/wav" | "audio/x-wav" => "wav",
"audio/mpeg" => "mp3",
"audio/ogg" => "ogg",
"video/mp4" => "mp4",
"video/webm" => "webm",
"application/pdf" => "pdf",
"application/json" => "json",
"text/plain" => "txt",
_ => "bin",
}
}
/// Inverse of [`ext_for_mime`] for serving persisted media with the right
/// `Content-Type`; generic binary for unknown extensions.
pub fn content_type_for_ext(ext: &str) -> &'static str {
match ext {
"png" => "image/png",
"jpg" => "image/jpeg",
"gif" => "image/gif",
"webp" => "image/webp",
"svg" => "image/svg+xml",
"wav" => "audio/wav",
"mp3" => "audio/mpeg",
"ogg" => "audio/ogg",
"mp4" => "video/mp4",
"webm" => "video/webm",
"pdf" => "application/pdf",
"json" => "application/json",
"txt" => "text/plain",
_ => "application/octet-stream",
}
}
+101
View File
@@ -0,0 +1,101 @@
//! Memory abstraction layer.
//!
//! Provides a [`Memory`] trait for pluggable long-term memory backends, and a
//! [`MemoryManager`] that holds at most **one** active backend at a time.
//!
//! # Singleton rule
//! Only one backend can be registered. If a second backend (with a different id)
//! tries to register, it is rejected with an `error!` log and the first one is
//! kept. The same backend can re-register itself (e.g. after a config change /
//! restart) — that replaces the existing registration cleanly.
//!
//! # Integration points
//! - [`Memory::query_context`] is called at the start of every `handle_message`
//! turn. The returned string is prepended to `extra_system_context` and
//! injected into the system prompt.
//! - [`Memory::tools`] is called per turn; the returned tools are added to the
//! LLM's tool list and dispatched before the global registry.
use std::sync::Arc;
use serde_json::Value;
use tokio::sync::RwLock;
use tracing::{error, info};
pub use core_api::memory::Memory;
use crate::tools::Tool;
// ── MemoryManager ─────────────────────────────────────────────────────────────
pub struct MemoryManager {
backend: RwLock<Option<Arc<dyn Memory>>>,
}
impl MemoryManager {
pub fn new() -> Self {
Self { backend: RwLock::new(None) }
}
/// Registers a memory backend.
///
/// - If no backend is registered yet, the new one is accepted.
/// - If the same backend id re-registers (restart / config change), it replaces
/// the old entry.
/// - If a **different** backend id tries to register while one is already active,
/// it is rejected with `error!` and the existing backend is kept.
pub async fn register(&self, backend: Arc<dyn Memory>) {
let mut lock = self.backend.write().await;
match lock.as_ref() {
None => {
info!("MemoryManager: registered backend '{}'", backend.id());
*lock = Some(backend);
}
Some(existing) if existing.id() == backend.id() => {
info!("MemoryManager: replacing backend '{}' (restart/reload)", backend.id());
*lock = Some(backend);
}
Some(existing) => {
error!(
"MemoryManager: backend '{}' is already registered — \
discarding '{}'. Only one memory backend is supported at a time.",
existing.id(),
backend.id(),
);
}
}
}
/// Returns memory context to inject into the system prompt for the upcoming
/// turn. Returns `None` if no backend is registered or the backend is
/// unavailable / has nothing to say.
pub async fn query_context(&self, session_id: i64, user_message: &str) -> Option<String> {
let backend = self.backend.read().await.clone()?;
if !backend.is_available() {
return None;
}
backend.query_context(session_id, user_message).await
}
/// Returns the per-turn LLM tools exposed by the active backend.
/// Empty if no backend is registered or the backend is unavailable.
pub async fn tools(&self) -> Vec<Arc<dyn Tool>> {
let backend = self.backend.read().await.clone();
match backend {
Some(b) if b.is_available() => b.tools(),
_ => vec![],
}
}
/// Builds OpenAI-format tool definitions from the active backend's tools.
pub async fn tool_defs(&self) -> Vec<Value> {
self.tools().await
.iter()
.map(|t| t.openai_definition())
.collect()
}
}
impl Default for MemoryManager {
fn default() -> Self { Self::new() }
}
+28
View File
@@ -0,0 +1,28 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
/// A structured notification produced by a background agent (TIC) or the cron
/// runner and delivered to the user's home conversation through `ChatHub`.
///
/// This replaces the previous free-text `String` briefing. Carrying `source`,
/// `event_type`, `event_time` and an open `refs` bag preserves the structured
/// context that already exists in `mcp_events` all the way to the main agent —
/// which is then the sole party responsible for the user-facing wording. The
/// `summary` is a neutral, third-person statement of fact, not a message
/// addressed to the user.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Notification {
/// Origin of the event: `"gmail"` | `"whatsapp"` | `"gcal"` | `"cron"` | `"system"`.
pub source: String,
/// Event kind, e.g. `"new_email"` | `"whatsapp_message"` | `"cron_result"`.
pub event_type: String,
/// Neutral, third-person factual summary of the event. NOT a message to the
/// user — the main agent phrases the user-facing message from this.
pub summary: String,
/// ISO 8601 timestamp of the underlying event.
pub event_time: String,
/// Open bag of actionable references (`message_id`, `thread_id`, `from`, …).
/// Defaults to an empty object when absent, keeping the shape forward-compatible.
#[serde(default)]
pub refs: Value,
}
+140
View File
@@ -0,0 +1,140 @@
//! Generic in-memory registry for pending human-in-the-loop requests.
//!
//! Approval, clarification, and elicitation all share the same shape: a request
//! is registered under an id with some display `Info`, a caller blocks on a
//! `oneshot::Receiver<Resolution>`, and later something resolves the request by
//! id — firing the sender and dropping the entry. This type factors out that
//! shared plumbing (the `Mutex<HashMap>` + oneshot bookkeeping) so each manager
//! keeps only what is genuinely its own: id minting, event emission, and any
//! extra policy (rules/bypass for approval, secret handling for elicitation).
//!
//! What deliberately stays OUT of the registry:
//! - **id minting** — the caller supplies the key (a durable `tool_call_id` for
//! approval, an internal counter for clarification/elicitation);
//! - **event emission** — the `ServerEvent` variants differ per manager, so the
//! caller broadcasts after `insert` / `resolve`;
//! - **ordering** — `list()` is unsorted; callers that need a stable order sort
//! on their own `Info` field (e.g. `created_at`).
use std::collections::HashMap;
use tokio::sync::{Mutex, oneshot};
/// One registered request: its display `Info` and the sender that unblocks the
/// waiting caller with a `Resolution`.
struct Entry<I, R> {
info: I,
tx: oneshot::Sender<R>,
}
/// Keyed store of pending requests. `I` is the cloneable public info surfaced to
/// the Inbox; `R` is the resolution payload delivered back to the blocked caller.
pub struct PendingRegistry<I, R> {
pending: Mutex<HashMap<i64, Entry<I, R>>>,
}
impl<I: Clone, R> PendingRegistry<I, R> {
pub fn new() -> Self {
Self { pending: Mutex::new(HashMap::new()) }
}
/// Registers `info` under `id` and returns the receiver the caller awaits.
/// The caller mints `id` (a durable tool_call_id or an internal counter).
pub async fn insert(&self, id: i64, info: I) -> oneshot::Receiver<R> {
let (tx, rx) = oneshot::channel();
self.pending.lock().await.insert(id, Entry { info, tx });
rx
}
/// Removes the entry for `id` and delivers `resolution` to the waiting caller.
/// Returns the entry's `info` (so the caller can broadcast a resolved event),
/// or `None` when no live entry exists (already resolved, or post-restart).
pub async fn resolve(&self, id: i64, resolution: R) -> Option<I> {
let entry = self.pending.lock().await.remove(&id)?;
let _ = entry.tx.send(resolution);
Some(entry.info)
}
/// Removes the entry for `id` WITHOUT sending a resolution: the dropped sender
/// makes the blocked caller observe `RecvError`. Used for deadline / disconnect
/// cancellation. Returns the removed `info`, or `None` if absent.
pub async fn remove(&self, id: i64) -> Option<I> {
self.pending.lock().await.remove(&id).map(|e| e.info)
}
/// Snapshot of the `info` for a single pending id, without resolving it.
pub async fn get(&self, id: i64) -> Option<I> {
self.pending.lock().await.get(&id).map(|e| e.info.clone())
}
/// Snapshot of every pending `info`, in unspecified order.
pub async fn list(&self) -> Vec<I> {
self.pending.lock().await.values().map(|e| e.info.clone()).collect()
}
/// Drops every entry whose `info` matches `pred` (their senders are dropped, so
/// the blocked callers observe `RecvError`). Returns the number removed.
pub async fn remove_where(&self, pred: impl Fn(&I) -> bool) -> usize {
let mut map = self.pending.lock().await;
let before = map.len();
map.retain(|_, e| !pred(&e.info));
before - map.len()
}
}
impl<I: Clone, R> Default for PendingRegistry<I, R> {
fn default() -> Self { Self::new() }
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn insert_then_resolve_delivers_and_returns_info() {
let reg: PendingRegistry<i64, String> = PendingRegistry::new();
let rx = reg.insert(42, 100).await;
// resolve returns the stored info and unblocks the waiter with the payload.
assert_eq!(reg.resolve(42, "answer".to_string()).await, Some(100));
assert_eq!(rx.await.unwrap(), "answer");
// the entry is gone afterwards.
assert!(reg.get(42).await.is_none());
}
#[tokio::test]
async fn resolve_unknown_id_is_none() {
let reg: PendingRegistry<i64, String> = PendingRegistry::new();
assert_eq!(reg.resolve(1, "x".to_string()).await, None);
}
#[tokio::test]
async fn remove_drops_sender_so_receiver_errors() {
let reg: PendingRegistry<i64, String> = PendingRegistry::new();
let rx = reg.insert(7, 100).await;
assert_eq!(reg.remove(7).await, Some(100));
// no resolution was sent — the dropped sender makes the waiter observe RecvError.
assert!(rx.await.is_err());
}
#[tokio::test]
async fn get_and_list_reflect_pending() {
let reg: PendingRegistry<i64, String> = PendingRegistry::new();
let _rx1 = reg.insert(1, 10).await;
let _rx2 = reg.insert(2, 20).await;
assert_eq!(reg.get(1).await, Some(10));
let mut all = reg.list().await;
all.sort();
assert_eq!(all, vec![10, 20]);
}
#[tokio::test]
async fn remove_where_filters_and_counts() {
let reg: PendingRegistry<i64, String> = PendingRegistry::new();
let _rx_keep = reg.insert(1, 10).await;
let rx_drop = reg.insert(2, 20).await;
// remove every entry whose info is >= 20.
assert_eq!(reg.remove_where(|info| *info >= 20).await, 1);
assert!(rx_drop.await.is_err()); // dropped entry's waiter errors
assert_eq!(reg.get(1).await, Some(10)); // the other entry stays
}
}
+351
View File
@@ -0,0 +1,351 @@
// The manager only ever handles `Arc<dyn Plugin>`. Naming a concrete plugin crate
// here would make the core depend on every plugin — and, through
// `plugin-transcribe-whisper-local`, on a C build — for no gain: the consumer
// constructs the plugin list and passes it to `Skald::new`.
pub use core_api::plugin::{Plugin, PluginContext, RouterFactory};
use std::collections::HashMap;
use std::sync::{Arc, OnceLock};
use std::time::Duration;
const PLUGIN_START_TIMEOUT_SECS: u64 = 30;
const PLUGIN_STOP_TIMEOUT_SECS: u64 = 5;
use anyhow::Result;
use serde::Serialize;
use serde_json::{Value, json};
use sqlx::SqlitePool;
use tokio::sync::Mutex;
use tokio::time::timeout;
use tracing::{error, info, warn};
use crate::db::plugins as db;
use crate::skald::Skald;
// ── Public plugin info (returned by list_items tool and REST API) ─────────────
#[derive(Debug, Clone, Serialize)]
pub struct PluginInfo {
pub id: String,
pub name: String,
pub description: String,
pub enabled: bool,
pub running: bool,
pub config: Value,
pub config_schema: Value,
pub runtime_status: Option<Value>,
}
// ── PluginManager ─────────────────────────────────────────────────────────────
pub struct PluginManager {
plugins: Vec<Arc<dyn Plugin>>,
db: Arc<SqlitePool>,
skald: OnceLock<Arc<Skald>>,
/// Provided by WebFrontend before start_enabled() is called.
router_factory: OnceLock<RouterFactory>,
/// HTTP port the web server is bound to — provided by WebFrontend before start_enabled().
web_port: OnceLock<u16>,
/// Last known (enabled, config_json) per plugin id — used by the watcher.
known_state: Mutex<HashMap<String, (bool, String)>>,
}
impl PluginManager {
pub fn new(db: Arc<SqlitePool>) -> Self {
Self {
plugins: Vec::new(),
db,
skald: OnceLock::new(),
router_factory: OnceLock::new(),
web_port: OnceLock::new(),
known_state: Mutex::new(HashMap::new()),
}
}
pub fn register(&mut self, plugin: impl Plugin + 'static) {
self.plugins.push(Arc::new(plugin));
}
pub fn register_arc(&mut self, plugin: Arc<dyn Plugin>) {
self.plugins.push(plugin);
}
pub fn set_skald(&self, skald: Arc<Skald>) {
let _ = self.skald.set(skald);
}
/// Called by WebFrontend before start_enabled().
pub fn set_router_factory(&self, factory: RouterFactory) {
let _ = self.router_factory.set(factory);
}
/// Called by WebFrontend before start_enabled().
pub fn set_web_port(&self, port: u16) {
let _ = self.web_port.set(port);
}
fn skald(&self) -> Result<Arc<Skald>> {
self.skald.get().cloned()
.ok_or_else(|| anyhow::anyhow!("PluginManager: skald not initialized"))
}
fn build_context(&self, skald: &Skald) -> Result<PluginContext> {
let router_factory = self.router_factory.get().cloned()
.ok_or_else(|| anyhow::anyhow!("PluginManager: router_factory not set"))?;
let web_port = self.web_port.get().copied()
.ok_or_else(|| anyhow::anyhow!("PluginManager: web_port not set"))?;
Ok(PluginContext {
chat_hub: Arc::clone(skald.chat_hub()) as _,
command: Arc::clone(skald.command_manager()) as _,
approval: Arc::clone(skald.approval()) as _,
inbox: Arc::new(skald.inbox().clone()) as _,
db: Arc::clone(skald.db()),
secrets: Arc::clone(skald.secrets()) as _,
transcribe: Arc::clone(skald.transcribe_manager()) as _,
transcribe_registry: Arc::clone(skald.transcribe_manager()) as _,
image_generate_registry: Arc::clone(skald.image_generator_manager()) as _,
tts_registry: Arc::clone(skald.tts_manager()) as _,
tts_provider: Arc::clone(skald.tts_manager()) as _,
api_provider_registry: Arc::clone(skald.provider_registry()) as _,
location: Arc::clone(skald.location_manager()) as _,
event_bus: Arc::clone(skald.event_bus()),
system_bus: Arc::clone(skald.system_bus()),
web_port,
remote_slot: Arc::clone(skald.remote()),
router_factory,
})
}
/// Collects the HTTP routers contributed by enabled plugins (plugin.md §12.3).
/// Returns `(plugin_id, router)` pairs; the caller (`WebFrontend::start`)
/// nests each under `/api/plugin/<id>/`. Only plugins with `enabled=true` in
/// the DB and a non-`None` `http_router()` are included.
///
/// Call this AFTER `start_enabled()` so a plugin's router can close over state
/// initialised during `reload`/`start`.
pub async fn collect_plugin_routers(&self) -> Vec<(String, axum::Router)> {
let mut out = Vec::new();
for plugin in &self.plugins {
match db::get(&self.db, plugin.id()).await {
Ok(Some(row)) if row.enabled => {}
Ok(_) => continue,
Err(e) => {
warn!(plugin = plugin.id(), error = %e, "collect_plugin_routers: DB read failed; skipping");
continue;
}
}
if let Some(router) = plugin.http_router() {
info!(plugin = plugin.id(), "plugin contributed an HTTP router → /api/plugin/{}", plugin.id());
out.push((plugin.id().to_string(), router));
}
}
out
}
// ── Startup ───────────────────────────────────────────────────────────────
/// Calls reload() for every plugin that has enabled=true in DB.
/// Plugins without a DB row are skipped (not yet configured).
/// After each successful start, registers the plugin's Memory backend (if any).
/// Must be called after both set_skald() and set_router_factory().
pub async fn start_enabled(&self) -> Result<()> {
let skald = self.skald()?;
// Build a full inventory for the bootstrap report: active (started),
// failed (enabled but errored/timed out), and available (disabled).
let mut active: Vec<String> = Vec::new();
let mut failed: Vec<(String, String)> = Vec::new();
let mut disabled: Vec<String> = Vec::new();
for plugin in &self.plugins {
let row = db::get(&self.db, plugin.id()).await?;
let enabled = row.as_ref().map(|r| r.enabled).unwrap_or(false);
if !enabled {
disabled.push(plugin.id().to_string());
continue;
}
let row = row.expect("enabled implies row present");
let config = serde_json::from_str(&row.config).unwrap_or(json!({}));
let deadline = Duration::from_secs(PLUGIN_START_TIMEOUT_SECS);
let ctx = self.build_context(&skald)?;
match timeout(deadline, plugin.reload(true, config, ctx)).await {
Ok(Ok(())) => {
self.known_state.lock().await
.insert(plugin.id().to_string(), (true, row.config));
info!(plugin = plugin.id(), "plugin started");
if let Some(mem) = plugin.memory() {
skald.memory_manager().register(mem).await;
}
active.push(plugin.id().to_string());
}
Ok(Err(e)) => {
error!(plugin = plugin.id(), error = %e, "plugin failed to start");
failed.push((plugin.id().to_string(), e.to_string()));
}
Err(_) => {
error!(plugin = plugin.id(), secs = PLUGIN_START_TIMEOUT_SECS, "plugin start timed out");
failed.push((plugin.id().to_string(),
format!("start timed out after {PLUGIN_START_TIMEOUT_SECS}s")));
}
}
}
crate::boot::section(format!(
"Plugins — {} active, {} failed, {} available",
active.len(), failed.len(), disabled.len()
));
if !active.is_empty() {
crate::boot::ok(active.join(", "));
}
for (id, reason) in &failed {
crate::boot::fail(format!("{id}{reason}"));
}
if !disabled.is_empty() {
crate::boot::off(disabled.join(", "));
}
Ok(())
}
pub async fn stop_all(&self) {
for plugin in &self.plugins {
if plugin.is_running() {
let deadline = Duration::from_secs(PLUGIN_STOP_TIMEOUT_SECS);
match timeout(deadline, plugin.stop()).await {
Ok(Ok(())) => info!(plugin = plugin.id(), "plugin stopped"),
Ok(Err(e)) => error!(plugin = plugin.id(), error = %e, "plugin stop error"),
Err(_) => warn!(plugin = plugin.id(), secs = PLUGIN_STOP_TIMEOUT_SECS, "plugin stop timed out"),
}
}
}
}
// ── Config update (called by REST API) ────────────────────────────────────
/// Persists the new config to DB, then calls reload() immediately.
pub async fn update_config(&self, id: &str, enabled: bool, config: Value) -> Result<()> {
let plugin = self.find(id)?;
let config_json = serde_json::to_string(&config)?;
db::upsert(&self.db, id, enabled, &config_json).await?;
let skald = self.skald()?;
plugin.reload(enabled, config, self.build_context(&skald)?).await?;
self.known_state.lock().await
.insert(id.to_string(), (enabled, config_json));
info!(plugin = id, enabled, "plugin config updated");
Ok(())
}
/// Toggle only the enabled flag, keeping existing config.
pub async fn toggle(&self, id: &str, enabled: bool) -> Result<()> {
let row = db::get(&self.db, id).await?
.unwrap_or_else(|| crate::db::plugins::PluginRow {
id: id.to_string(),
enabled,
config: "{}".to_string(),
});
let config: Value = serde_json::from_str(&row.config).unwrap_or(json!({}));
self.update_config(id, enabled, config).await
}
// ── Background config watcher ─────────────────────────────────────────────
/// Spawns a Tokio task that polls the DB every 30 s and calls reload()
/// on any plugin whose (enabled, config) has changed since last check.
/// This is the fallback path; normal updates go through update_config().
pub fn start_config_watcher(self: &Arc<Self>, shutdown: tokio_util::sync::CancellationToken) {
let this = Arc::clone(self);
tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_secs(30));
interval.tick().await; // skip immediate first tick
loop {
tokio::select! {
_ = shutdown.cancelled() => { break; }
_ = interval.tick() => {
if let Err(e) = this.check_and_reload().await {
error!(error = %e, "plugin config watcher error");
}
}
}
}
});
}
async fn check_and_reload(&self) -> Result<()> {
let rows = db::list(&self.db).await?;
let skald = self.skald()?;
// Collect what needs reloading while holding the lock briefly.
let to_reload: Vec<_> = {
let known = self.known_state.lock().await;
rows.into_iter()
.filter(|row| {
known.get(&row.id)
.map_or(true, |(e, c)| *e != row.enabled || c != &row.config)
})
.collect()
};
for row in to_reload {
let Ok(plugin) = self.find(&row.id) else { continue };
let config = serde_json::from_str(&row.config).unwrap_or(json!({}));
let ctx = self.build_context(&skald)?;
match plugin.reload(row.enabled, config, ctx).await {
Ok(()) => {
self.known_state.lock().await
.insert(row.id.clone(), (row.enabled, row.config));
info!(plugin = row.id, "plugin reloaded by config watcher");
if row.enabled {
if let Some(mem) = plugin.memory() {
skald.memory_manager().register(mem).await;
}
}
}
Err(e) => error!(plugin = row.id, error = %e, "plugin reload failed"),
}
}
Ok(())
}
// ── Queries ───────────────────────────────────────────────────────────────
pub async fn list(&self) -> Result<Vec<PluginInfo>> {
let mut out = Vec::new();
for plugin in &self.plugins {
let row = db::get(&self.db, plugin.id()).await?;
let (enabled, config_json) = row
.map(|r| (r.enabled, r.config))
.unwrap_or((false, "{}".to_string()));
out.push(PluginInfo {
id: plugin.id().to_string(),
name: plugin.name().to_string(),
description: plugin.description().to_string(),
enabled,
running: plugin.is_running(),
config: serde_json::from_str(&config_json).unwrap_or(json!({})),
config_schema: plugin.config_schema(),
runtime_status: plugin.runtime_status(),
});
}
Ok(out)
}
/// Every registered plugin, enabled or not. Lets the core ask each one for
/// its contributions (`Plugin::tools`, `Plugin::http_router`) without ever
/// naming a concrete plugin type.
pub fn all(&self) -> &[Arc<dyn Plugin>] {
&self.plugins
}
pub fn get_plugin_typed<T: Plugin + 'static>(&self, id: &str) -> Option<Arc<T>> {
self.plugins.iter()
.find(|p| p.id() == id)
.and_then(|p| Arc::clone(p).as_arc_any().downcast::<T>().ok())
}
fn find(&self, id: &str) -> Result<Arc<dyn Plugin>> {
self.plugins.iter()
.find(|p| p.id() == id)
.cloned()
.ok_or_else(|| anyhow::anyhow!("plugin not found: {id}"))
}
}
+108
View File
@@ -0,0 +1,108 @@
pub mod tickets;
use std::sync::Arc;
use anyhow::Result;
use sqlx::SqlitePool;
use crate::db::projects::{self, Project};
use crate::run_context::RunContext;
pub struct ProjectManager {
db: Arc<SqlitePool>,
}
impl ProjectManager {
pub fn new(db: Arc<SqlitePool>) -> Self {
Self { db }
}
pub async fn list(&self) -> Result<Vec<Project>> {
projects::list(&self.db).await
}
pub async fn get(&self, id: i64) -> Result<Option<Project>> {
projects::get(&self.db, id).await
}
pub async fn create(
&self,
name: &str,
path: &str,
description: &str,
run_context: Option<&RunContext>,
) -> Result<Project> {
let rc_json = run_context.map(|rc| rc.to_db());
projects::create(&self.db, name, path, description, rc_json.as_deref()).await
}
pub async fn update(
&self,
id: i64,
name: &str,
path: &str,
description: &str,
run_context: Option<&RunContext>,
) -> Result<bool> {
let rc_json = run_context.map(|rc| rc.to_db());
projects::update(&self.db, id, name, path, description, rc_json.as_deref()).await
}
pub async fn delete(&self, id: i64) -> Result<bool> {
projects::delete(&self.db, id).await
}
}
/// Builds the runtime `RunContext` for working on `project`, layering project-runtime
/// fields over an optional pre-resolved `base` RC (which carries static config set at
/// creation time, e.g. `security_group`).
///
/// Runtime fields computed here:
/// - `working_directory` — always set to `project.path`.
/// - `allow_fs_writes` — project tree + Skald's own `data/` directory.
/// - `system_prompt` — project-context fragments prepended before any stored ones.
///
/// Shared by `ProjectTicketManager::start` (background ticket jobs) and the interactive
/// project-chat session provisioning, so both work with identical context.
pub fn build_runtime_run_context(project: &Project, base: Option<RunContext>) -> RunContext {
let mut rc = base.unwrap_or_default();
// Working directory is always the project path, overwritten at build time.
rc.working_directory = Some(project.path.clone());
// Absolute path to Skald's own data directory (user personal data store).
let skald_data = std::env::current_dir()
.unwrap_or_default()
.join("data")
.to_string_lossy()
.into_owned();
// Grant write access to the project tree and Skald's data directory.
if !rc.allow_fs_writes.contains(&project.path) {
rc.allow_fs_writes.push(project.path.clone());
}
if !rc.allow_fs_writes.contains(&skald_data) {
rc.allow_fs_writes.push(skald_data.clone());
}
// Build runtime context fragments and prepend before any stored ones.
// Note: working directory is intentionally omitted here — the date/time/OS/WD
// tail block in MessageBuilder already reflects the effective WD from RunContext.
let project_header = if project.description.is_empty() {
format!("You are working on project \"{}\".", project.name)
} else {
format!("You are working on project \"{}\". Description: {}", project.name, project.description)
};
let mut injected = vec![
project_header,
format!(
"Personal user data is available at: {}. \
Consult it when the task requires knowledge about the user.",
skald_data
),
];
injected.extend(std::mem::take(&mut rc.system_prompt));
rc.system_prompt = injected;
rc
}
+178
View File
@@ -0,0 +1,178 @@
use std::sync::Arc;
use anyhow::{Result, anyhow};
use sqlx::SqlitePool;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use core_api::system_bus::{SystemEvent, SystemEventBus};
use crate::cron::TaskManager;
use crate::db::{project_tickets, project_tickets::ProjectTicket, projects};
use crate::run_context::RunContext;
pub struct ProjectTicketManager {
db: Arc<SqlitePool>,
task_mgr: std::sync::OnceLock<Arc<TaskManager>>,
}
impl ProjectTicketManager {
pub fn new(db: Arc<SqlitePool>) -> Arc<Self> {
Arc::new(Self {
db,
task_mgr: std::sync::OnceLock::new(),
})
}
pub fn set_task_manager(&self, tm: Arc<TaskManager>) {
let _ = self.task_mgr.set(tm);
}
/// Subscribe to the system bus and react to `JobCompleted` events whose
/// `origin_ref` starts with `"PROJECT_TASK:"`. Spawns a background task.
pub fn start_listener(
self: Arc<Self>,
system_bus: Arc<SystemEventBus>,
shutdown: CancellationToken,
) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
let mut rx = system_bus.subscribe();
loop {
tokio::select! {
_ = shutdown.cancelled() => break,
res = rx.recv() => {
match res {
Ok(SystemEvent::JobCompleted { origin_ref: Some(ref s), result, error, .. })
if s.starts_with("PROJECT_TASK:") =>
{
if let Some(tid) = s.strip_prefix("PROJECT_TASK:")
.and_then(|n| n.parse::<i64>().ok())
{
if let Err(e) = self.on_job_completed(
tid,
result.as_deref(),
error.as_deref(),
).await {
warn!(error = %e, ticket_id = tid, "ticket completion failed");
}
}
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
warn!("ProjectTicketManager: system_bus lagged by {n} events");
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
_ => {}
}
}
}
}
})
}
// ── CRUD ─────────────────────────────────────────────────────────────────
pub async fn list(&self, project_id: i64) -> Result<Vec<ProjectTicket>> {
project_tickets::list_for_project(&self.db, project_id).await
}
pub async fn get(&self, id: i64) -> Result<Option<ProjectTicket>> {
project_tickets::get(&self.db, id).await
}
pub async fn create(
&self,
project_id: i64,
title: &str,
description: &str,
agent_id: &str,
run_context: Option<&RunContext>,
) -> Result<ProjectTicket> {
let rc_json = run_context.map(|rc| rc.to_db());
let ticket = project_tickets::create(
&self.db, project_id, title, description, agent_id, rc_json.as_deref(),
).await?;
projects::touch(&self.db, project_id).await?;
Ok(ticket)
}
pub async fn delete(&self, id: i64) -> Result<bool> {
let ticket = project_tickets::get(&self.db, id).await?;
let found = project_tickets::delete(&self.db, id).await?;
if found {
if let Some(t) = ticket {
projects::touch(&self.db, t.project_id).await?;
}
}
Ok(found)
}
// ── Lifecycle ─────────────────────────────────────────────────────────────
/// Builds a runtime RunContext and starts the ticket as a background job.
///
/// The stored RC (ticket → project) carries only static config set at creation
/// time (e.g. `security_group`). All runtime fields are computed here:
/// - `working_directory` — always set to `project.path`
/// - `allow_fs_writes` — project tree + Skald's own `data/` directory
/// - `system_prompt` — project context fragments prepended before any stored ones
pub async fn start(&self, ticket_id: i64) -> Result<()> {
let task_mgr = self.task_mgr.get()
.ok_or_else(|| anyhow!("ProjectTicketManager: task_manager not initialized"))?;
let ticket = project_tickets::get(&self.db, ticket_id).await?
.ok_or_else(|| anyhow!("ticket {ticket_id} not found"))?;
let project = projects::get(&self.db, ticket.project_id).await?
.ok_or_else(|| anyhow!("project {} not found", ticket.project_id))?;
// Resolve base RC (ticket override → project default → empty), then layer the
// project-runtime fields (WD, fs-write grants, project-context system prompt).
// The stored RC carries only static config (e.g. security_group set at creation).
let base: Option<RunContext> =
ticket.run_context.as_deref().and_then(RunContext::from_db)
.or_else(|| project.run_context.as_deref().and_then(RunContext::from_db));
let rc = super::build_runtime_run_context(&project, base);
let origin_ref = format!("PROJECT_TASK:{ticket_id}");
let rc_json = rc.to_db();
let job = task_mgr.spawn_async_job(
&ticket.title,
&ticket.description,
&ticket.description,
&ticket.agent_id,
Some(&rc_json),
&origin_ref,
)?;
project_tickets::start(&self.db, ticket_id, job.id).await?;
projects::touch(&self.db, ticket.project_id).await?;
Ok(())
}
/// Called when a `SystemEvent::JobCompleted` with matching `origin_ref` is received.
async fn on_job_completed(
&self,
ticket_id: i64,
result: Option<&str>,
error: Option<&str>,
) -> Result<()> {
let project_id = project_tickets::get(&self.db, ticket_id).await?
.map(|t| t.project_id);
project_tickets::complete(&self.db, ticket_id, result, error).await?;
if let Some(pid) = project_id {
projects::touch(&self.db, pid).await?;
}
Ok(())
}
/// Reset a ticket back to todo, clearing all run state.
pub async fn reset(&self, ticket_id: i64) -> Result<()> {
let project_id = project_tickets::get(&self.db, ticket_id).await?
.map(|t| t.project_id);
project_tickets::reset(&self.db, ticket_id).await?;
if let Some(pid) = project_id {
projects::touch(&self.db, pid).await?;
}
Ok(())
}
}
+77
View File
@@ -0,0 +1,77 @@
// All provider types and traits now live in core-api.
// Re-export everything so existing imports in this crate continue to work.
pub use core_api::provider::{
ApiProvider, ApiProviderRegistry, BuiltLlmClient,
LlmModelRecord, LlmProviderRecord, LlmStrength,
ProviderField, ProviderUiMeta, ReasoningMode, ServiceType,
RemoteLlmModelInfo,
};
// ── ProviderRegistry ──────────────────────────────────────────────────────────
use std::sync::Arc;
use core_api::system_bus::{SystemEvent, SystemEventBus};
pub struct ProviderRegistry {
builtin: Vec<Arc<dyn ApiProvider>>,
plugins: std::sync::RwLock<Vec<Arc<dyn ApiProvider>>>,
system_bus: Arc<SystemEventBus>,
}
impl ProviderRegistry {
pub fn new(system_bus: Arc<SystemEventBus>) -> Self {
Self {
builtin: Vec::new(),
plugins: std::sync::RwLock::new(Vec::new()),
system_bus,
}
}
pub fn register_builtin(&mut self, p: impl ApiProvider + 'static) {
self.builtin.push(Arc::new(p));
}
/// Looks up a provider by type_id. Plugin providers shadow built-in ones.
pub fn get(&self, type_id: &str) -> Option<Arc<dyn ApiProvider>> {
{
let plugins = self.plugins.read().unwrap();
if let Some(p) = plugins.iter().find(|p| p.type_id() == type_id) {
return Some(Arc::clone(p));
}
}
self.builtin.iter().find(|p| p.type_id() == type_id).cloned()
}
/// Returns all known providers: plugin-registered first, then built-in.
pub fn all(&self) -> Vec<Arc<dyn ApiProvider>> {
let plugins = self.plugins.read().unwrap();
let mut result: Vec<Arc<dyn ApiProvider>> = plugins.clone();
for p in &self.builtin {
if !result.iter().any(|x| x.type_id() == p.type_id()) {
result.push(Arc::clone(p));
}
}
result
}
pub fn contains(&self, type_id: &str) -> bool {
self.get(type_id).is_some()
}
}
impl core_api::provider::ApiProviderRegistry for ProviderRegistry {
fn register_plugin(&self, p: Arc<dyn ApiProvider>) {
let id = p.type_id();
let mut plugins = self.plugins.write().unwrap();
plugins.retain(|x| x.type_id() != id);
plugins.push(p);
tracing::info!(type_id = id, "provider registered (plugin)");
self.system_bus.send(SystemEvent::ApiProviderRegistered { type_id: id.to_string() });
}
fn unregister_plugin(&self, type_id: &str) {
self.plugins.write().unwrap().retain(|p| p.type_id() != type_id);
tracing::info!(type_id, "provider unregistered (plugin)");
self.system_bus.send(SystemEvent::ApiProviderUnregistered { type_id: type_id.to_string() });
}
}
+376
View File
@@ -0,0 +1,376 @@
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::{Result, bail};
use serde::{Deserialize, Serialize};
use sqlx::SqlitePool;
use tracing::info;
pub use crate::db::tool_permission_groups::ToolPermissionGroup;
use crate::approval::{ApprovalManager, RuleAction};
use crate::tools::fs::{canonicalize_for_policy, path_under};
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct RunContext {
security_group: Option<String>,
#[serde(default)]
pub system_prompt: Vec<String>,
#[serde(default)]
pub allow_fs_writes: Vec<String>,
/// Extra directories/files granted read-only access (beyond the working directory,
/// `docs/`, `skills/`, and everything in `allow_fs_writes`, which is readable too).
#[serde(default)]
pub allow_fs_reads: Vec<String>,
/// Working directory for tool calls. None means Skald's own process cwd.
#[serde(default)]
pub working_directory: Option<String>,
}
impl RunContext {
pub fn with_security_group(security_group: Option<String>) -> Self {
Self { security_group, ..Default::default() }
}
pub fn to_db(&self) -> String {
serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
}
pub fn from_db(s: &str) -> Option<Self> {
if s.is_empty() { return None; }
serde_json::from_str(s).ok()
}
/// Permission group ID for approval rule lookup.
pub fn tool_group_id(&self) -> Option<&str> {
self.security_group.as_deref()
}
/// Combined system prompt fragments to inject as dynamic context, or None if empty.
pub fn extra_system_prompt(&self) -> Option<String> {
if self.system_prompt.is_empty() { return None; }
Some(self.system_prompt.join("\n\n"))
}
/// Effective working directory for this session.
/// Returns the configured path if set and non-empty, otherwise Skald's process cwd.
pub fn effective_working_dir(&self) -> PathBuf {
self.working_directory
.as_deref()
.filter(|d| !d.is_empty())
.map(PathBuf::from)
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default())
}
/// True if writing to `path` is pre-authorized by this RunContext.
/// Entries in `allow_fs_writes` are resolved against `effective_working_dir`,
/// so relative entries like `"data"` are treated as relative to the session WD.
/// Paths are canonicalized first (resolving `..`/symlinks), then matched as
/// exact file OR recursive directory prefix.
pub fn is_write_allowed(&self, path: &str) -> bool {
if self.allow_fs_writes.is_empty() { return false; }
let wd = self.effective_working_dir();
let canon = canonicalize_for_policy(path, &wd);
self.allow_fs_writes.iter().any(|entry| {
path_under(&canon, &canonicalize_for_policy(entry, &wd))
})
}
/// True if reading `path` is pre-authorized by this RunContext.
/// Read access is granted (no approval prompt) for: the working directory itself,
/// its `docs/` and `skills/` subtrees (always-safe baseline), any `allow_fs_reads`
/// entry, and anything writable (write implies read). All paths are canonicalized
/// first so `..`/symlink escapes cannot widen the grant.
///
/// Note: this only relaxes a `Require` decision to `Allow` — an explicit `Deny`
/// rule (e.g. on `secrets/`) still wins, because the approval engine is consulted
/// first and `Deny` is never overridden by this fast-path.
pub fn is_read_allowed(&self, path: &str) -> bool {
let wd = self.effective_working_dir();
let canon = canonicalize_for_policy(path, &wd);
let mut roots: Vec<std::path::PathBuf> = vec![
canonicalize_for_policy(".", &wd), // working directory itself
canonicalize_for_policy("docs", &wd),
canonicalize_for_policy("skills", &wd),
];
roots.extend(self.allow_fs_reads.iter().map(|e| canonicalize_for_policy(e, &wd)));
roots.extend(self.allow_fs_writes.iter().map(|e| canonicalize_for_policy(e, &wd)));
roots.iter().any(|root| path_under(&canon, root))
}
}
pub struct RunContextManager {
db: Arc<SqlitePool>,
approval: Arc<ApprovalManager>,
}
impl RunContextManager {
pub fn new(db: Arc<SqlitePool>, approval: Arc<ApprovalManager>) -> Self {
Self { db, approval }
}
/// Seeds the built-in "default" permission group and migrates legacy rules.
/// Safe to call at every startup (idempotent).
pub async fn seed_defaults(&self) -> Result<()> {
crate::db::tool_permission_groups::insert_or_ignore(
&self.db, "default", "Default", Some("Built-in default permission group"),
).await?;
let migrated = sqlx::query("UPDATE approval_rules SET group_id = 'default' WHERE group_id IS NULL")
.execute(self.db.as_ref())
.await
.map(|r| r.rows_affected())
.unwrap_or(0);
if migrated > 0 {
info!(%migrated, "run_context: migrated approval rules to 'default' group");
}
Ok(())
}
// ── ToolPermissionGroup CRUD ───────────────────────────────────────────────
pub async fn list_groups(&self) -> Result<Vec<ToolPermissionGroup>> {
crate::db::tool_permission_groups::list(&self.db).await
}
pub async fn get_group(&self, id: &str) -> Result<Option<ToolPermissionGroup>> {
crate::db::tool_permission_groups::get(&self.db, id).await
}
pub async fn create_group(
&self,
id: &str,
name: &str,
description: Option<&str>,
) -> Result<()> {
if id == "default" {
bail!("cannot create a permission group with reserved id 'default'");
}
crate::db::tool_permission_groups::insert(&self.db, id, name, description).await
}
pub async fn update_group(
&self,
id: &str,
name: &str,
description: Option<&str>,
) -> Result<bool> {
crate::db::tool_permission_groups::update(&self.db, id, name, description).await
}
pub async fn delete_group(&self, id: &str) -> Result<bool> {
if id == "default" {
bail!("cannot delete the built-in 'default' permission group");
}
crate::db::tool_permission_groups::delete(&self.db, id).await
}
/// Duplicates a permission group and all its rules atomically.
pub async fn duplicate_group(
&self,
source_id: &str,
new_id: &str,
new_name: &str,
) -> Result<()> {
if new_id == "default" {
bail!("cannot create a permission group with reserved id 'default'");
}
let source = crate::db::tool_permission_groups::get(&self.db, source_id).await?
.ok_or_else(|| anyhow::anyhow!("source group '{source_id}' not found"))?;
let mut tx = self.db.begin().await?;
sqlx::query(
"INSERT INTO tool_permission_groups (id, name, description) VALUES (?, ?, ?)",
)
.bind(new_id)
.bind(new_name)
.bind(source.description.as_deref())
.execute(&mut *tx)
.await?;
sqlx::query(
"INSERT INTO approval_rules \
(agent_id, source, tool_pattern, path_pattern, action, note, priority, group_id) \
SELECT agent_id, source, tool_pattern, path_pattern, action, note, priority, ? \
FROM approval_rules \
WHERE group_id = ?",
)
.bind(new_id)
.bind(source_id)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
// ── Tool visibility ────────────────────────────────────────────────────────
/// Returns the effective `RuleAction` for `tool_name` under the given permission group.
/// `run_context_id` now directly holds a `tool_permission_groups` id (the run_contexts
/// table indirection has been removed). Falls back to the `"default"` group when `None`.
pub async fn check_tool_visibility(
&self,
run_context_id: Option<&str>,
tool_name: &str,
) -> Option<RuleAction> {
let group_id = run_context_id.unwrap_or("default");
self.approval.check_tool_visibility(group_id, tool_name).await
}
// ── Session assignment ─────────────────────────────────────────────────────
/// Serialises `ctx` as JSON and stores it on the session row.
/// `None` clears the context (falls back to the default permission group).
pub async fn set_session_run_context(
&self,
session_id: i64,
ctx: Option<&RunContext>,
) -> Result<()> {
let json = ctx.map(|rc| rc.to_db());
sqlx::query("UPDATE chat_sessions SET run_context = ? WHERE id = ?")
.bind(json.as_deref())
.bind(session_id)
.execute(self.db.as_ref())
.await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
/// Creates a fresh, uniquely-named temp directory for an fs test.
fn unique_tmp() -> PathBuf {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
let dir = std::env::temp_dir()
.join(format!("skald_rc_test_{}_{}", std::process::id(), nanos));
std::fs::create_dir_all(&dir).unwrap();
dir
}
fn rc_with_wd(wd: &PathBuf) -> RunContext {
RunContext {
working_directory: Some(wd.to_string_lossy().into_owned()),
..Default::default()
}
}
#[test]
fn read_allows_working_dir_docs_skills() {
let wd = unique_tmp();
for sub in ["docs", "skills", "sub", "secrets"] {
std::fs::create_dir_all(wd.join(sub)).unwrap();
std::fs::write(wd.join(sub).join("f.txt"), "x").unwrap();
}
std::fs::write(wd.join("root.txt"), "x").unwrap();
let rc = rc_with_wd(&wd);
assert!(rc.is_read_allowed("root.txt"));
assert!(rc.is_read_allowed("docs/f.txt"));
assert!(rc.is_read_allowed("skills/f.txt"));
assert!(rc.is_read_allowed("sub/f.txt"));
// secrets/ is under the WD, so the fast-path allows it — the `secrets/` *deny rule*
// (consulted before this fast-path in the gate) is what actually blocks it.
assert!(rc.is_read_allowed("secrets/f.txt"));
std::fs::remove_dir_all(&wd).ok();
}
#[test]
fn read_denies_outside_working_dir() {
let wd = unique_tmp();
let outside = unique_tmp(); // sibling temp dir, not under wd
std::fs::write(outside.join("f.txt"), "x").unwrap();
let rc = rc_with_wd(&wd);
assert!(!rc.is_read_allowed(outside.join("f.txt").to_str().unwrap()));
std::fs::remove_dir_all(&wd).ok();
std::fs::remove_dir_all(&outside).ok();
}
#[test]
fn read_allows_write_paths_and_extra_reads() {
let wd = unique_tmp();
let writable = unique_tmp();
let readable = unique_tmp();
std::fs::write(writable.join("w.txt"), "x").unwrap();
std::fs::write(readable.join("r.txt"), "x").unwrap();
let rc = RunContext {
working_directory: Some(wd.to_string_lossy().into_owned()),
allow_fs_writes: vec![writable.to_string_lossy().into_owned()],
allow_fs_reads: vec![readable.to_string_lossy().into_owned()],
..Default::default()
};
// write implies read
assert!(rc.is_read_allowed(writable.join("w.txt").to_str().unwrap()));
assert!(rc.is_write_allowed(writable.join("w.txt").to_str().unwrap()));
// read-only grant: readable but not writable
assert!(rc.is_read_allowed(readable.join("r.txt").to_str().unwrap()));
assert!(!rc.is_write_allowed(readable.join("r.txt").to_str().unwrap()));
std::fs::remove_dir_all(&wd).ok();
std::fs::remove_dir_all(&writable).ok();
std::fs::remove_dir_all(&readable).ok();
}
#[test]
fn canonicalize_resolves_parent_traversal() {
let wd = unique_tmp();
std::fs::create_dir_all(wd.join("docs")).unwrap();
std::fs::create_dir_all(wd.join("secrets")).unwrap();
std::fs::write(wd.join("secrets").join("s.txt"), "x").unwrap();
assert_eq!(
canonicalize_for_policy("docs/../secrets/s.txt", &wd),
canonicalize_for_policy("secrets/s.txt", &wd),
);
std::fs::remove_dir_all(&wd).ok();
}
#[test]
fn canonicalize_resolves_symlink_escape() {
let wd = unique_tmp();
std::fs::create_dir_all(wd.join("docs")).unwrap();
std::fs::create_dir_all(wd.join("secrets")).unwrap();
std::fs::write(wd.join("secrets").join("s.txt"), "x").unwrap();
std::os::unix::fs::symlink(wd.join("secrets"), wd.join("docs").join("leak")).unwrap();
// A symlink docs/leak -> secrets must resolve to the real secrets path.
assert_eq!(
canonicalize_for_policy("docs/leak/s.txt", &wd),
canonicalize_for_policy("secrets/s.txt", &wd),
);
std::fs::remove_dir_all(&wd).ok();
}
#[test]
fn write_allow_not_bypassed_by_traversal() {
let wd = unique_tmp();
std::fs::create_dir_all(wd.join("data")).unwrap();
std::fs::create_dir_all(wd.join("secrets")).unwrap();
let rc = RunContext {
working_directory: Some(wd.to_string_lossy().into_owned()),
allow_fs_writes: vec!["data".to_string()],
..Default::default()
};
// Writing into data/ is allowed...
assert!(rc.is_write_allowed("data/new.txt"));
// ...but data/../secrets/x escapes the grant and must NOT be allowed.
assert!(!rc.is_write_allowed("data/../secrets/x.txt"));
std::fs::remove_dir_all(&wd).ok();
}
}
+65
View File
@@ -0,0 +1,65 @@
use std::sync::Arc;
use anyhow::Result;
use async_trait::async_trait;
use sqlx::SqlitePool;
use tracing::debug;
pub use core_api::secrets::{SecretsApi, require};
// ── SecretsStore ──────────────────────────────────────────────────────────────
pub struct SecretsStore {
pool: Arc<SqlitePool>,
}
impl SecretsStore {
pub fn new(pool: Arc<SqlitePool>) -> Arc<Self> {
Arc::new(Self { pool })
}
}
#[async_trait]
impl SecretsApi for SecretsStore {
async fn get(&self, key: &str) -> Option<String> {
sqlx::query_scalar::<_, String>(
"SELECT value FROM secrets WHERE key = ?1",
)
.bind(key)
.fetch_optional(&*self.pool)
.await
.ok()
.flatten()
}
async fn set(&self, key: &str, value: &str) -> Result<()> {
sqlx::query(
"INSERT INTO secrets (key, value)
VALUES (?1, ?2)
ON CONFLICT(key) DO UPDATE SET value = excluded.value,
updated_at = datetime('now')",
)
.bind(key)
.bind(value)
.execute(&*self.pool)
.await?;
debug!(key, "secret set");
Ok(())
}
async fn delete(&self, key: &str) -> Result<()> {
sqlx::query("DELETE FROM secrets WHERE key = ?1")
.bind(key)
.execute(&*self.pool)
.await?;
debug!(key, "secret deleted");
Ok(())
}
async fn list_keys(&self) -> Vec<String> {
sqlx::query_scalar::<_, String>("SELECT key FROM secrets ORDER BY key ASC")
.fetch_all(&*self.pool)
.await
.unwrap_or_default()
}
}
+9
View File
@@ -0,0 +1,9 @@
use crate::provider::ServiceType;
/// Light umbrella trait shared by all model managers.
/// Enables grouping managers generically (e.g. models-hub routing, diagnostics)
/// without coupling to their specific CRUD operations.
pub trait ServiceManager: Send + Sync {
fn service_type(&self) -> ServiceType;
fn display_name(&self) -> &'static str;
}
@@ -0,0 +1,329 @@
use std::collections::HashSet;
use std::sync::{Arc, RwLock};
use serde_json::Value;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tracing::info;
use crate::db::{chat_history, chat_llm_tools, chat_sessions_stack, scratchpad, stack_mcp_grants};
use crate::events::ServerEvent;
use super::{ChatSessionHandler, MAX_AGENT_DEPTH, TurnOutcome};
use super::emitter::TurnEmitter;
use super::interface_tools::{AgentRunConfig, InterfaceTool, ToolFuture};
use super::config::activate_tools_tool_def;
impl ChatSessionHandler {
/// Dispatches a sub-agent as a child stack frame within the current session.
/// Used by `execute_task` (mode=sync) and `execute_subtask` interceptions in `llm_loop`.
/// Args must contain `agent_id` and `prompt`; optionally `client`.
pub(super) async fn dispatch_sub_agent(
&self,
parent_stack_id: i64,
parent_config: &AgentRunConfig,
parent_tool_call_id: i64,
args: &Value,
token: &CancellationToken,
tx: &mpsc::Sender<ServerEvent>,
) -> anyhow::Result<String> {
let pool = &self.db;
let em = TurnEmitter::new(tx);
let target_id = args["agent_id"].as_str()
.ok_or_else(|| anyhow::anyhow!("dispatch_sub_agent: missing required argument `agent_id`"))?;
let prompt = args["prompt"].as_str()
.ok_or_else(|| anyhow::anyhow!("dispatch_sub_agent: missing required argument `prompt`"))?;
if target_id == parent_config.agent_id {
anyhow::bail!("dispatch_sub_agent: an agent cannot call itself (`{target_id}`)");
}
// Only `task` agents are dispatchable: this rejects `chat` (e.g. `main`,
// `project-coordinator`) and `system` (e.g. `tic`) agents, and surfaces a
// not-found error for unknown ids — all in one gate.
let target_meta = crate::agents::load_task_meta(target_id)
.map_err(|e| anyhow::anyhow!("dispatch_sub_agent: {e}"))?;
let parent_frame = chat_sessions_stack::find_by_id(pool, parent_stack_id).await?
.ok_or_else(|| anyhow::anyhow!("dispatch_sub_agent: parent stack frame not found"))?;
let new_depth = parent_frame.depth + 1;
if new_depth > MAX_AGENT_DEPTH {
anyhow::bail!(
"dispatch_sub_agent: maximum agent depth ({}) exceeded — refusing to recurse further",
MAX_AGENT_DEPTH
);
}
let explicit_client = args["client"].as_str().or(target_meta.client.as_deref());
let (resolved_client, _) = self.llm_manager.resolve(
explicit_client,
target_meta.scope.as_deref(),
target_meta.strength,
).await.map_err(|e| anyhow::anyhow!("dispatch_sub_agent: {e}"))?;
let child = chat_sessions_stack::create(
pool,
self.session_id,
target_id,
Some(prompt),
new_depth,
Some(parent_tool_call_id),
).await?;
let persisted_grants = stack_mcp_grants::list_for_stack(pool, child.id)
.await
.unwrap_or_default();
let active_mcp_grants: Arc<RwLock<HashSet<String>>> =
Arc::new(RwLock::new(persisted_grants.into_iter().collect()));
let mut child_config = parent_config.for_sub_agent(target_id.to_string(), resolved_client.clone());
child_config.active_mcp_grants = Arc::clone(&active_mcp_grants);
child_config.base_tool_defs.extend(self.tools.openai_definitions_sub_agents_only());
child_config.base_tool_defs.push(super::ask_user_clarification_tool_def());
// Let the sub-agent dispatch a further sub-agent (e.g. tech-lead → architect/engineer).
// `execute_subtask` is intercepted in `run_agent_turn` and routed back here. Only expose it
// while the child can still recurse — at the depth limit `dispatch_sub_agent` would reject it.
if new_depth < MAX_AGENT_DEPTH {
child_config.base_tool_defs.push(super::execute_subtask_tool_def());
}
{
let group_id = self.tool_group_id().await;
let gid = group_id.as_deref().unwrap_or("default");
let group_rules = crate::db::approval_rules::list_for_group(
pool, Some(gid),
).await.unwrap_or_default();
child_config.base_tool_defs.retain(|def| {
let name = def["function"]["name"].as_str().unwrap_or("");
self.approval.is_tool_visible(&group_rules, name)
});
}
{
let pool_clone = Arc::clone(&self.db);
let session_id = self.session_id;
let stack_id = child.id;
let mcp_clone = Arc::clone(&self.mcp);
let grants_clone = Arc::clone(&active_mcp_grants);
let activate_tool = crate::tools::activate_tools::ActivateTools {
pool: pool_clone,
session_id,
stack_id: Some(stack_id),
mcp: mcp_clone,
active_mcp_grants: grants_clone,
};
let activate_tool = Arc::new(activate_tool);
child_config.interface_tools.push(InterfaceTool {
definition: activate_tools_tool_def(),
handler: Arc::new(move |args| -> ToolFuture {
use crate::tools::Tool as _;
let tool = Arc::clone(&activate_tool);
Box::pin(async move {
tokio::task::spawn_blocking(move || tool.execute(args))
.await
.map_err(|e| anyhow::anyhow!("activate_tools task panicked: {e}"))?
})
}),
});
}
chat_history::append(pool, child.id, &chat_history::Role::Agent, prompt, false, None).await?;
let prompt_preview = super::preview_truncate(prompt, 500);
em.agent_start(
child.id,
parent_tool_call_id,
target_id.to_string(),
parent_config.agent_id.clone(),
new_depth,
prompt_preview,
).await;
info!(
session_id = self.session_id,
parent_stack = parent_stack_id,
child_stack = child.id,
target_agent = target_id,
client = %resolved_client,
"dispatch_sub_agent: running child inline"
);
// Run the child synchronously in the SAME task, holding the same
// `processing` lock and sharing the same cancellation token. The returned
// string becomes the parent tool call's result, which `run_agent_turn`
// persists and emits as `ToolDone` — so completion lives in one place.
// Boxed: `resume_pending_tools` now dispatches sub-agents via `execute_tool_call`,
// which re-enters here — box this edge so the recursive async future stays sized.
let _ = Box::pin(self.resume_pending_tools(child.id, &child_config, token, tx)).await;
// Sub-agents never inject live user input.
let outcome = self.run_agent_turn(child.id, &child_config, token, tx, None).await;
if let Err(e) = stack_mcp_grants::delete_for_stack(pool, child.id).await {
tracing::warn!(stack_id = child.id, error = %e, "dispatch_sub_agent: failed to delete stack MCP grants");
}
let parent_agent_id = parent_config.agent_id.clone();
let child_agent_id = target_id.to_string();
let preview = |s: &str| super::preview_truncate(s, 500);
let result = match outcome {
Ok(TurnOutcome::Final { content, .. }) => {
em.agent_done(child.id, child_agent_id, parent_agent_id, preview(&content)).await;
Ok(content)
}
Ok(TurnOutcome::Cancelled) => {
// The parent shares this token: if the cancel came from the user,
// its next round check returns Cancelled too. We still record a
// tool result so the history stays well-formed.
em.agent_done(child.id, child_agent_id, parent_agent_id, "⚠️ Cancelled.".to_string()).await;
Ok(format!("Sub-agent `{target_id}` was cancelled."))
}
Ok(TurnOutcome::Exhausted) => {
em.agent_done(child.id, child_agent_id, parent_agent_id, "⚠️ Exhausted tool-call rounds.".to_string()).await;
Ok(format!(
"Sub-agent `{target_id}` exceeded {} tool-call rounds without producing a final answer.",
self.max_tool_rounds
))
}
Err(e) => {
let msg = e.to_string();
em.agent_done(child.id, child_agent_id, parent_agent_id, format!("⚠️ Error: {msg}")).await;
Err(e)
}
};
let _ = chat_sessions_stack::terminate(pool, child.id).await;
result
}
/// Handles the `update_scratchpad` built-in.
///
/// The scratchpad is a session-scoped shared blackboard (`scratchpad_sid()` is
/// the session_id, identical for every frame). When a homogeneous batch of
/// sub-agents runs concurrently (`handle_sub_agent_batch`), two siblings writing
/// the *same* key race to last-writer-wins — this is inherent to a shared
/// blackboard and accepted by design, not a correctness bug. Sub-agents that must
/// not clobber each other should write distinct keys.
pub(super) async fn dispatch_update_scratchpad(
&self,
args: &Value,
) -> anyhow::Result<String> {
let key = args["key"].as_str().unwrap_or("").to_string();
let value = args["value"].as_str().unwrap_or("").to_string();
scratchpad::upsert(&self.db, self.scratchpad_sid(), &key, &value).await
.map(|_| format!("Scratchpad updated: {key}"))
}
/// Handles the `write_todos` built-in.
///
/// Stateless: the list is not persisted anywhere — it lives only in this
/// agent's tool-result history (per-stack, so it is never seen by sub-agents
/// or the caller). We just validate/normalise the items and echo back a
/// formatted checklist the model re-reads from its own tool result.
pub(super) async fn dispatch_write_todos(
&self,
args: &Value,
) -> anyhow::Result<String> {
let items = args["todos"].as_array().ok_or_else(|| {
anyhow::anyhow!("`write_todos` requires a `todos` array. Re-send the full list, e.g. [{{\"content\":\"...\",\"status\":\"pending\"}}].")
})?;
if items.is_empty() {
return Err(anyhow::anyhow!("`todos` is empty — send at least one item, or omit the call entirely."));
}
let mut lines = Vec::with_capacity(items.len());
let (mut done, mut active, mut pending) = (0usize, 0usize, 0usize);
for item in items {
let content = item["content"].as_str().unwrap_or("").trim();
if content.is_empty() {
continue;
}
// Normalise unknown statuses to `pending`.
let marker = match item["status"].as_str() {
Some("completed") => { done += 1; "x" }
Some("in_progress") => { active += 1; "~" }
_ => { pending += 1; " " }
};
lines.push(format!("[{marker}] {content}"));
}
if lines.is_empty() {
return Err(anyhow::anyhow!("No valid todo items (every `content` was empty)."));
}
Ok(format!(
"Todo list ({total}): {done} done, {active} in progress, {pending} pending\n{body}",
total = lines.len(),
body = lines.join("\n"),
))
}
/// Handles the `ask_user_clarification` built-in.
///
/// Interactive sessions (web, telegram): sends `AgentQuestion` over the WS channel
/// and waits for the user to answer inline in the chat.
///
/// Background sessions (cron, tic): registers in `ClarificationManager` so the
/// Agent Inbox page can surface and resolve the request.
///
/// `tool_call_id` is used to mark the DB row as `pending` before blocking,
/// so page refreshes and app restarts can distinguish "waiting for input" from
/// "was executing" and re-ask the question correctly.
pub(super) async fn dispatch_ask_user_clarification(
&self,
tool_call_id: i64,
args: &Value,
tx: &mpsc::Sender<ServerEvent>,
) -> anyhow::Result<String> {
let title = args["title"].as_str().unwrap_or("Clarification needed").to_string();
let question = args["question"].as_str().unwrap_or("?").to_string();
let suggested: Vec<String> = args["suggested_answers"]
.as_array()
.map(|a| a.iter().filter_map(|v| v.as_str().map(str::to_string)).collect())
.unwrap_or_default();
// Mark as pending before suspending so restart/refresh can re-ask the question.
chat_llm_tools::set_approval_pending(&self.db, tool_call_id).await?;
let context_label = self.context_label.read().ok().and_then(|g| g.clone());
// Always register in ClarificationManager so the question appears in the
// Agent Inbox for ALL sessions (both interactive web/telegram and background cron/tic).
let (request_id, rx) = self.clarification.register(
self.session_id,
&self.agent_id,
&self.source,
context_label.as_deref(),
&title,
&question,
suggested.clone(),
).await;
tracing::debug!(session_id = self.session_id, request_id, is_interactive = self.is_interactive, source = %self.source, "dispatch_ask_user_clarification: routing");
if self.is_interactive {
// For interactive sessions, also send the question over WS so it appears
// inline in the chat. The user can answer from either the chat or the Inbox.
info!(session_id = self.session_id, request_id, %question, source = %self.source, "agent asking user for clarification (interactive) — sending AgentQuestion");
let send_result = tx.send(ServerEvent::AgentQuestion {
request_id,
tool_call_id,
title,
question,
suggested_answers: suggested,
}).await;
if send_result.is_err() {
tracing::warn!(session_id = self.session_id, request_id, "AgentQuestion send failed — tx receiver dropped");
} else {
info!(session_id = self.session_id, request_id, "AgentQuestion sent to bridge");
}
} else {
info!(session_id = self.session_id, request_id, %question, source = %self.source, "background session waiting for clarification");
}
// Wait for the answer (from WS via resolve_question → clarification.resolve,
// or directly from the Inbox REST endpoint).
rx.await.map_err(|_| anyhow::Error::new(super::AgentFlowSignal::QuestionChannelClosed))
}
}
@@ -0,0 +1,115 @@
use serde_json::Value;
use tracing::debug;
use super::ChatSessionHandler;
use super::emitter::TurnEmitter;
use crate::tools::{is_file_write_tool, tool_names as tn};
impl ChatSessionHandler {
/// Emits the appropriate frontend approval event for the given tool call.
///
/// | Tool kind | Event emitted |
/// |------------------|-------------------------------------------------------|
/// | file-write tools | `PendingWrite` with before/after diff (IO concurrent) |
/// | `execute_cmd` | `PendingWrite` with command preview |
/// | `restart` | `PendingWrite` with restart description |
/// | everything else | `ApprovalRequired` |
///
/// Called from both `llm_loop` and `resume_pending_tools` to avoid duplication.
pub(super) async fn emit_approval_event(
&self,
em: &TurnEmitter<'_>,
request_id: i64,
tool_call_id: i64,
tool_name: &str,
arguments: &Value,
) {
if is_file_write_tool(tool_name) {
let path = arguments["path"].as_str().unwrap_or("").to_string();
// Read current file and compute new content concurrently — both are disk I/O.
let (old_content, new_content) = tokio::join!(
self.read_current_content(&path),
self.compute_new_content(tool_name, arguments),
);
if let Some(new_content) = new_content {
em.pending_write(request_id, tool_call_id, path, old_content, new_content).await;
} else {
// File doesn't exist yet or diff can't be computed — fall back to generic.
debug!(tool = tool_name, "emit_approval_event: no diff available, using ApprovalRequired");
em.approval_required(request_id, tool_call_id, tool_name.to_string(), arguments.clone()).await;
}
} else if tool_name == tn::EXECUTE_CMD {
let cmd = arguments["command"].as_str().unwrap_or("");
em.pending_write(request_id, tool_call_id, "$ execute_cmd".to_string(), None, format!("$ {cmd}")).await;
} else if tool_name == tn::RESTART {
em.pending_write(
request_id, tool_call_id,
"$ restart".to_string(),
None,
"Riavvia il processo (exit -1 → supervisor ricompila e rilancia)".to_string(),
).await;
} else {
em.approval_required(request_id, tool_call_id, tool_name.to_string(), arguments.clone()).await;
}
}
/// Reads the current content of a file from disk (for diff generation in PendingWrite events).
pub(super) async fn read_current_content(&self, path: &str) -> Option<String> {
let abs = crate::tools::fs::resolve(path).ok()?;
tokio::fs::read_to_string(&abs).await.ok()
}
/// Computes what a file would look like after the tool runs, without writing it.
/// Returns `None` if the result cannot be determined (e.g. edit_file on a missing file).
pub(super) async fn compute_new_content(&self, name: &str, args: &Value) -> Option<String> {
match name {
"write_file" => args["content"].as_str().map(|s| s.to_string()),
"edit_file" => {
let path = args["path"].as_str()?;
let old_text = args["old"].as_str()?;
let new_text = args["new"].as_str()?;
let current = self.read_current_content(path).await?;
if current.contains(old_text) {
Some(current.replacen(old_text, new_text, 1))
} else {
None
}
}
"insert_at_line" => {
let path = args["path"].as_str()?;
let line_num = args["line"].as_u64()? as usize;
let new_text = args["content"].as_str()?;
let placement = args["placement"].as_str().unwrap_or("after");
if line_num == 0 { return None; }
let current = self.read_current_content(path).await?;
let mut lines: Vec<&str> = current.split('\n').collect();
let idx = (line_num - 1).min(lines.len().saturating_sub(1));
let insert_idx = if placement == "before" { idx } else { idx + 1 };
let new_lines: Vec<&str> = new_text.split('\n').collect();
for (i, l) in new_lines.iter().enumerate() {
lines.insert(insert_idx + i, l);
}
Some(lines.join("\n"))
}
"replace_lines" => {
let path = args["path"].as_str()?;
let from_line = args["from_line"].as_u64()? as usize;
let to_line = args["to_line"].as_u64()? as usize;
let new_text = args["new"].as_str()?;
if from_line == 0 || to_line < from_line { return None; }
let current = self.read_current_content(path).await?;
let mut lines: Vec<&str> = current.lines().collect();
let total = lines.len();
if from_line > total { return None; }
let to_clamped = to_line.min(total);
let new_lines: Vec<&str> = new_text.lines().collect();
lines.splice((from_line - 1)..to_clamped, new_lines);
let has_trailing = current.ends_with('\n');
let mut result = lines.join("\n");
if has_trailing { result.push('\n'); }
Some(result)
}
_ => None,
}
}
}
@@ -0,0 +1,205 @@
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};
use serde_json::Value;
use crate::tools::tool_names as tn;
use super::{ChatSessionHandler, update_scratchpad_tool_def, write_todos_tool_def};
use super::interface_tools::{AgentRunConfig, InterfaceTool, ToolFuture};
/// Returns an `activate_tools` OpenAI tool definition.
pub(super) fn activate_tools_tool_def() -> Value {
serde_json::json!({
"type": "function",
"function": {
"name": tn::ACTIVATE_TOOLS,
"description": "Activate one or more tool groups so their tools become available. \
A group is either an MCP server name (see the MCP list) or the reserved \
keyword `config`, which loads all system-configuration tools (managing \
MCP servers, plugins, scheduled cron jobs, and secrets). \
Pass an array of group names (e.g. [\"gmail\", \"config\"]). \
Once activated, the tools are available from the next tool-call round onward.",
"parameters": {
"type": "object",
"properties": {
"groups": {
"type": "array",
"items": { "type": "string" },
"description": "Tool groups to activate: MCP server names and/or the reserved \
keyword \"config\" (e.g. [\"gmail\", \"config\"])."
}
},
"required": ["groups"]
}
}
})
}
impl ChatSessionHandler {
/// Resolves the LLM client and assembles `AgentRunConfig` for a top-level turn
/// (depth = 0). Extracted to avoid duplicating the same ~15 lines in both
/// `handle_message` and `resume_turn`.
pub(super) async fn build_agent_config(
&self,
client_name: Option<String>,
extra_system: Option<String>,
extra_system_dynamic: Option<String>,
mut interface_tools: Vec<InterfaceTool>,
system_substitutions: HashMap<String, String>,
) -> anyhow::Result<AgentRunConfig> {
let meta = crate::agents::load_meta(&self.agent_id).ok();
let (key, _) = self.llm_manager.resolve(
client_name.as_deref(),
meta.as_ref().and_then(|m| m.scope.as_deref()),
meta.as_ref().and_then(|m| m.strength),
).await?;
let mut base_tool_defs = self.tools.openai_definitions_excluding_config();
// Config-category built-ins are hidden from the always-on set and lazy-loaded
// via `activate_tools(["config"])`. They go through the same interactive-only /
// approval-visibility filters as base_tool_defs below, then ride in AgentRunConfig
// as `config_tool_defs` (appended by `all_tool_defs()` only when granted).
let mut config_tool_defs = self.tools.openai_definitions_config_only();
base_tool_defs.push(update_scratchpad_tool_def());
base_tool_defs.push(write_todos_tool_def());
// `ask_user_clarification` is available to every agent except hidden `system`
// agents (e.g. TIC), which have no user-facing channel. Interactive sessions
// emit AgentQuestion inline (plus the Inbox); background sessions rely on the
// Inbox alone.
let is_system = meta
.as_ref()
.map(|m| m.agent_type == crate::agents::AgentType::System)
.unwrap_or(false);
if !is_system {
base_tool_defs.push(super::ask_user_clarification_tool_def());
}
// Background sessions (cron, tic): remove tools that only make sense in
// interactive sessions (e.g. read_notification, which is synthetically
// injected by ChatHub and returns EMPTY if called directly).
if !self.is_interactive {
let interactive_only = self.tools.interactive_only_names();
let keep = |def: &Value| {
let name = def["function"]["name"].as_str().unwrap_or("");
!interactive_only.iter().any(|n| n == name)
};
base_tool_defs.retain(|d| keep(d));
config_tool_defs.retain(|d| keep(d));
}
// Interactive sessions get read_agent_result so the LLM can poll for async
// task status. The real delivery happens via inject_async_result (synthetic msg).
if self.is_interactive {
base_tool_defs.push(serde_json::json!({
"type": "function",
"function": {
"name": "task_completed",
"description": "Invoked BY THE SYSTEM (not by you) when an async task finishes, \
delivering its result. You will never need to call this yourself — \
the system calls it automatically when execute_task(mode=async) completes.",
"parameters": {
"type": "object",
"required": ["task_id"],
"properties": {
"task_id": { "type": "integer", "description": "The completed task id" }
}
}
}
}));
}
// Approval-rules visibility filter: hide tools whose effective action for
// this session's permission group is Deny. Rules are loaded once and applied
// synchronously; the execution-time gate in ApprovalManager remains as a
// second layer of enforcement.
{
let group_id = self.tool_group_id().await;
let gid = group_id.as_deref().unwrap_or("default");
let group_rules = crate::db::approval_rules::list_for_group(
&self.db, Some(gid),
).await.unwrap_or_default();
let visible = |def: &Value| {
let name = def["function"]["name"].as_str().unwrap_or("");
self.approval.is_tool_visible(&group_rules, name)
};
base_tool_defs.retain(|d| visible(d));
config_tool_defs.retain(|d| visible(d));
}
// ── Tool-group grant initialisation ─────────────────────────────────────
//
// Load persisted session grants from DB (MCP server names and/or the reserved
// `config` keyword), then inject `activate_tools` so the LLM can activate
// additional groups on demand.
let persisted = crate::db::session_mcp_grants::list_for_session(
&self.db, self.session_id,
).await.unwrap_or_default();
let active_mcp_grants: Arc<RwLock<HashSet<String>>> =
Arc::new(RwLock::new(persisted.into_iter().collect()));
{
let pool_clone = Arc::clone(&self.db);
let session_id = self.session_id;
let mcp_clone = Arc::clone(&self.mcp);
let grants_clone = Arc::clone(&active_mcp_grants);
let activate_tool = crate::tools::activate_tools::ActivateTools {
pool: pool_clone,
session_id,
stack_id: None,
mcp: mcp_clone,
active_mcp_grants: grants_clone,
};
let activate_tool = Arc::new(activate_tool);
interface_tools.push(InterfaceTool {
definition: activate_tools_tool_def(),
handler: Arc::new(move |args| -> ToolFuture {
use crate::tools::Tool as _;
let tool = Arc::clone(&activate_tool);
Box::pin(async move {
tokio::task::spawn_blocking(move || tool.execute(args))
.await
.map_err(|e| anyhow::anyhow!("activate_tools task panicked: {e}"))?
})
}),
});
}
// ── End tool-group grant initialisation ─────────────────────────────────
// Append RunContext system prompt fragments to the dynamic tail (not cached).
let extra_system_dynamic = {
let rc = self.run_context.read().await;
let injected = rc.as_ref().and_then(|r| r.extra_system_prompt());
match (extra_system_dynamic, injected) {
(Some(e), Some(i)) => Some(format!("{e}\n\n{i}")),
(Some(e), None) => Some(e),
(None, Some(i)) => Some(i),
(None, None) => None,
}
};
let root_only_tool_names: Vec<String> = self.tools.root_agent_only_names();
let memory_tools = self.memory_manager.tools().await;
let image_tools = Arc::clone(&self.image_generator_manager).tools().await;
Ok(AgentRunConfig {
agent_id: self.agent_id.clone(),
client_name: key,
depth: 0,
base_tool_defs,
config_tool_defs,
extra_system,
extra_system_dynamic,
tail_reminder: None,
system_substitutions,
interface_tools,
memory_tools,
image_tools,
mcp: Arc::clone(&self.mcp),
active_mcp_grants,
root_only_tool_names,
})
}
}
@@ -0,0 +1,149 @@
//! Working-directory argument rewriting and the per-tool-call dispatch router.
//!
//! Extracted from `run_agent_turn`: `effective_args` applies the RunContext working
//! directory to a call's arguments, and `execute_tool_call` routes an approved call
//! to the right executor (special non-cancellable paths + the unified cancellable
//! `ToolExecution` path).
use serde_json::Value;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use crate::events::ServerEvent;
use crate::tools::{drive_execution, tool_names as tn, ExecutionOutcome, ToolResult};
use super::ChatSessionHandler;
use super::interface_tools::AgentRunConfig;
/// Whether a tool call is a synchronous sub-agent dispatch, i.e. one intercepted
/// by `execute_tool_call` and routed to `dispatch_sub_agent` rather than the
/// registry. Covers `execute_task` (mode=sync), `execute_subtask`, and the legacy
/// `run_subtask` alias (only reachable via a `pending` call left across a restart).
/// Shared by the router below and the parallel-batch detection in `run_agent_turn`.
pub(super) fn is_sync_sub_agent(tool_name: &str, args: &Value) -> bool {
(tool_name == tn::EXECUTE_TASK && args["mode"].as_str() == Some("sync") && args.get("agent_id").is_some())
|| tool_name == tn::EXECUTE_SUBTASK
|| tool_name == "run_subtask"
}
/// Result of routing a single tool call to its executor.
pub(super) enum DispatchResult {
/// Normal completion / failure / cancellation — the caller records it.
Outcome(ExecutionOutcome),
/// The turn must end now and the tool row must stay `pending`: the
/// `ask_user_clarification` WS channel closed while awaiting an answer. The
/// caller returns `TurnOutcome::Cancelled` **without** recording the tool, so
/// `resume_pending_tools` re-asks it on reconnect.
AbortPending,
}
impl ChatSessionHandler {
/// Applies the RunContext working directory to a tool call's arguments:
/// resolves a relative `path` against the effective WD and injects `workdir`
/// for `execute_cmd`. The caller keeps the original `arguments` for the
/// `ToolStart` event / DB logging; this returns the copy used for execution.
pub(super) async fn effective_args(&self, tool_name: &str, args: &Value) -> Value {
let mut effective = args.clone();
let wd = self.run_context.read().await
.as_ref()
.map(|rc| rc.effective_working_dir());
if let Some(wd) = wd {
if let Some(path) = effective["path"].as_str()
&& !std::path::Path::new(path).is_absolute()
{
effective["path"] = Value::String(wd.join(path).to_string_lossy().into_owned());
}
if tool_name == tn::EXECUTE_CMD && effective.get("workdir").is_none() {
effective["workdir"] = Value::String(wd.to_string_lossy().into_owned());
}
}
effective
}
/// Routes one already-approved tool call to the right executor. Covers the
/// special, non-cancellable paths (sub-agent, scratchpad, todos, clarification,
/// the `task_completed` stub) and the unified cancellable `ToolExecution` path
/// (registry / memory / image / interface / MCP). `restart` is handled by the
/// caller before this is reached (it calls `_exit` and never returns).
#[allow(clippy::too_many_arguments)]
pub(super) async fn execute_tool_call(
&self,
stack_id: i64,
config: &AgentRunConfig,
tool_call_id: i64,
tool_name: &str,
args: &Value,
token: &CancellationToken,
tx: &mpsc::Sender<ServerEvent>,
) -> DispatchResult {
let outcome: ExecutionOutcome = if is_sync_sub_agent(tool_name, args) {
plain_outcome(self.dispatch_sub_agent(stack_id, config, tool_call_id, args, token, tx).await)
} else if tool_name == tn::UPDATE_SCRATCHPAD {
plain_outcome(self.dispatch_update_scratchpad(args).await)
} else if tool_name == tn::WRITE_TODOS {
plain_outcome(self.dispatch_write_todos(args).await)
} else if tool_name == tn::ASK_USER_CLARIFICATION {
match self.dispatch_ask_user_clarification(tool_call_id, args, tx).await {
Ok(answer) => ExecutionOutcome::Completed(ToolResult::Text(answer)),
Err(err) => {
// WS disconnected while waiting for a clarification answer.
// Tool stays 'pending' in DB — resume_pending_tools re-dispatches on reconnect.
if matches!(err.downcast_ref::<super::AgentFlowSignal>(), Some(super::AgentFlowSignal::QuestionChannelClosed)) {
warn!(session_id = self.session_id, tool_call_id, "clarification channel closed — aborting turn (tool stays pending)");
return DispatchResult::AbortPending;
}
ExecutionOutcome::Failed(err.to_string())
}
}
} else if tool_name == "task_completed" {
// Defensive stub: if the LLM somehow calls this itself, return a hint.
// Real delivery is via inject_async_result (synthetic message from the system).
let task_id = args["task_id"].as_i64().unwrap_or(0);
ExecutionOutcome::Completed(ToolResult::Text(format!(r#"{{"status":"not_ready","task_id":{task_id},"message":"This tool is invoked by the system, not by you. Do not call it again — the result will arrive automatically as a new message in this conversation."}}"#)))
} else {
// Unified cancellable path. The execution owns its in-flight state and
// its own stop(); on /stop the work future is dropped (aborting I/O /
// killing the child) and the tool is recorded as Cancelled, not Failed.
match self.build_execution(tool_name, args.clone(), config) {
Some(exec) => drive_execution(exec.as_ref(), token).await,
None => ExecutionOutcome::Failed(format!("Unknown tool: {tool_name}")),
}
};
DispatchResult::Outcome(outcome)
}
}
/// Maps a plain dispatch `Result<String>` to an [`ExecutionOutcome`]. Used by the
/// non-cancellable special paths (sub-agent, scratchpad, todos), which can only
/// complete or fail — never `Cancelled`.
fn plain_outcome(result: anyhow::Result<String>) -> ExecutionOutcome {
match result {
Ok(s) => ExecutionOutcome::Completed(ToolResult::Text(s)),
Err(e) => ExecutionOutcome::Failed(e.to_string()),
}
}
#[cfg(test)]
mod tests {
use super::is_sync_sub_agent;
use serde_json::json;
#[test]
fn recognises_sync_sub_agent_calls() {
assert!(is_sync_sub_agent("execute_task", &json!({"mode": "sync", "agent_id": "x"})));
assert!(is_sync_sub_agent("execute_subtask", &json!({})));
assert!(is_sync_sub_agent("run_subtask", &json!({}))); // legacy alias
}
#[test]
fn rejects_everything_else() {
// execute_task without mode=sync + agent_id is NOT a sync sub-agent.
assert!(!is_sync_sub_agent("execute_task", &json!({"mode": "async", "agent_id": "x"})));
assert!(!is_sync_sub_agent("execute_task", &json!({"mode": "sync"}))); // no agent_id
assert!(!is_sync_sub_agent("execute_task", &json!({})));
// Regular tools never qualify (they must keep the sequential path).
assert!(!is_sync_sub_agent("read_file", &json!({"path": "/x"})));
assert!(!is_sync_sub_agent("execute_cmd", &json!({"cmd": "ls"})));
}
}
@@ -0,0 +1,155 @@
//! Typed, fire-and-forget event seam for a running agent turn.
//!
//! Every event a turn produces used to be sent inline as
//! `tx.send(ServerEvent::X { .. }).await.ok()`, scattered across `llm_loop`,
//! `resume`, `agent_dispatch`, and `approval`. `TurnEmitter` wraps the per-turn
//! `mpsc::Sender<ServerEvent>` (which `ChatHub` bridges onto the global broadcast
//! bus) and exposes one semantic method per event, so the loop speaks in domain
//! terms (`emitter.tool_done(..)`) instead of constructing wire enums by hand.
//!
//! It is a zero-cost borrow wrapper: construct one at the top of a function that
//! emits and pass `&TurnEmitter` to any helper. This is also the single seam a
//! future event-bus / UI-vs-domain split would hook into.
use serde_json::Value;
use tokio::sync::mpsc;
use core_api::message_meta::Attachment;
use crate::events::ServerEvent;
/// Borrows the per-turn event sender and emits typed [`ServerEvent`]s.
pub(super) struct TurnEmitter<'a> {
tx: &'a mpsc::Sender<ServerEvent>,
}
impl<'a> TurnEmitter<'a> {
pub(super) fn new(tx: &'a mpsc::Sender<ServerEvent>) -> Self {
Self { tx }
}
/// Send an event, dropping it silently if the receiver is gone (the same
/// `.await.ok()` semantics every call site used before).
async fn emit(&self, event: ServerEvent) {
self.tx.send(event).await.ok();
}
// ── User / assistant turn events ────────────────────────────────────────
/// A user message row was persisted (telnet-style echo).
pub(super) async fn user_message(&self, message_id: i64, content: String, attachments: Vec<Attachment>) {
self.emit(ServerEvent::UserMessage { message_id, content, attachments }).await;
}
/// The assistant produced text alongside tool calls (reasoning before acting).
pub(super) async fn thinking(&self, message_id: i64, content: String, input_tokens: Option<u32>, output_tokens: Option<u32>) {
self.emit(ServerEvent::Thinking { message_id, content, input_tokens, output_tokens }).await;
}
/// The assistant response is complete.
pub(super) async fn done(&self, message_id: i64, stack_id: i64, content: String, input_tokens: Option<u32>, output_tokens: Option<u32>) {
self.emit(ServerEvent::Done { message_id, stack_id, content, input_tokens, output_tokens }).await;
}
/// The LLM was cut off by the token limit.
pub(super) async fn truncated(&self, output_tokens: Option<u32>) {
self.emit(ServerEvent::Truncated { output_tokens }).await;
}
/// A fatal error occurred processing the request.
pub(super) async fn error(&self, message: String) {
self.emit(ServerEvent::Error { message }).await;
}
// ── Tool-call lifecycle ─────────────────────────────────────────────────
#[allow(clippy::too_many_arguments)]
pub(super) async fn tool_start(
&self,
tool_call_id: i64,
message_id: i64,
name: String,
arguments: Value,
label_short: String,
label_full: String,
path: Option<String>,
) {
self.emit(ServerEvent::ToolStart {
tool_call_id, message_id, name, arguments, label_short, label_full, path,
}).await;
}
pub(super) async fn tool_done(&self, tool_call_id: i64, result: String, result_type: String) {
self.emit(ServerEvent::ToolDone { tool_call_id, result, result_type }).await;
}
pub(super) async fn tool_error(&self, tool_call_id: i64, error: String) {
self.emit(ServerEvent::ToolError { tool_call_id, error }).await;
}
pub(super) async fn tool_cancelled(&self, tool_call_id: i64) {
self.emit(ServerEvent::ToolCancelled { tool_call_id }).await;
}
pub(super) async fn tool_rejected(&self, tool_call_id: i64, reason: String) {
self.emit(ServerEvent::ToolRejected { tool_call_id, reason }).await;
}
/// A file-write tool completed; ask clients holding the file to reload.
pub(super) async fn file_changed(&self, path: String) {
self.emit(ServerEvent::FileChanged { path }).await;
}
// ── Approval / clarification prompts ────────────────────────────────────
#[allow(clippy::too_many_arguments)]
pub(super) async fn pending_write(
&self,
request_id: i64,
tool_call_id: i64,
path: String,
old_content: Option<String>,
new_content: String,
) {
self.emit(ServerEvent::PendingWrite { request_id, tool_call_id, path, old_content, new_content }).await;
}
pub(super) async fn approval_required(&self, request_id: i64, tool_call_id: i64, tool_name: String, arguments: Value) {
self.emit(ServerEvent::ApprovalRequired { request_id, tool_call_id, tool_name, arguments }).await;
}
// Note: `AgentQuestion` is emitted directly in `dispatch_ask_user_clarification`
// because that one site inspects the send Result for diagnostic logging — it is
// deliberately not wrapped here.
// ── Sub-agent stack frames ──────────────────────────────────────────────
#[allow(clippy::too_many_arguments)]
pub(super) async fn agent_start(
&self,
stack_id: i64,
parent_tool_call_id: i64,
agent_id: String,
parent_agent_id: String,
depth: i64,
prompt_preview: String,
) {
self.emit(ServerEvent::AgentStart {
stack_id, parent_tool_call_id, agent_id, parent_agent_id, depth, prompt_preview,
}).await;
}
pub(super) async fn agent_done(&self, stack_id: i64, agent_id: String, parent_agent_id: String, result_preview: String) {
self.emit(ServerEvent::AgentDone { stack_id, agent_id, parent_agent_id, result_preview }).await;
}
// ── LLM model fallback ──────────────────────────────────────────────────
pub(super) async fn model_fallback(&self, from: String, to: String, reason: String) {
self.emit(ServerEvent::ModelFallback { from, to, reason }).await;
}
pub(super) async fn llm_failed(&self, tried: Vec<String>, last_error: String) {
self.emit(ServerEvent::LlmFailed { tried, last_error }).await;
}
}
@@ -0,0 +1,138 @@
//! Shared approval gate for a single tool call.
//!
//! The decision + human-approval flow (approval-engine check, RunContext
//! fast-path, auto-deny, register + await) was duplicated in `run_agent_turn` and
//! `resume_pending_tools`, and had already drifted (only the live loop applied the
//! RunContext fast-path and the auto-deny short-circuit). `run_approval_gate` is the
//! single implementation both call, so the two paths gate identically.
use std::sync::atomic::Ordering;
use serde_json::Value;
use tracing::{info, warn};
use crate::approval::GateResult;
use crate::db::chat_llm_tools;
use crate::run_context::RunContext;
use crate::tools::{is_file_read_tool, is_file_write_tool};
use super::{ApprovalDecision, ChatSessionHandler};
use super::emitter::TurnEmitter;
/// Result of the approval gate for a single tool call.
pub(super) enum GateOutcome {
/// The tool may execute.
Proceed,
/// Denied by policy, auto-denied, or rejected by a human. The DB row has been
/// marked `rejected` and the `ToolRejected` event emitted — the caller just
/// skips the call.
Rejected,
/// The approval channel closed (WS disconnected) while awaiting a decision.
/// The caller must end the turn / resume.
ChannelClosed,
}
impl ChatSessionHandler {
/// Runs a tool call through the approval engine and, when human approval is
/// required, registers the request, emits the approval event, and awaits the
/// decision. Shared by `run_agent_turn` and `resume_pending_tools`.
pub(super) async fn run_approval_gate(
&self,
tool_call_id: i64,
tool_name: &str,
args: &Value,
agent_id: &str,
em: &TurnEmitter<'_>,
) -> anyhow::Result<GateOutcome> {
let pool = &self.db;
// Post-restart manual resolve: this exact tool_call was already approved by the
// user via a resolve endpoint, which then triggered this resume. There is no
// live oneshot to unblock, so skip re-gating (and re-prompting) and dispatch it.
if self.pre_approved.lock().unwrap().remove(&tool_call_id) {
info!(session_id = self.session_id, tool = %tool_name, tool_call_id, "approval: pre-approved (post-restart resolve) — skipping gate");
return Ok(GateOutcome::Proceed);
}
let category = self.tools.category_of(tool_name);
let group_id = self.tool_group_id().await;
// The approval engine decides first: an explicit Deny/Allow rule always wins.
let mut gate = self.approval.check(
self.session_id, category,
agent_id, &self.source, tool_name, args,
group_id.as_deref(),
).await;
// RunContext fast-path: relax `Require` to `Allow` for pre-authorized
// filesystem paths. It never overrides a `Deny` (same semantics as session
// bypass), so e.g. the `secrets/` deny rule holds even inside an auto-read
// working directory.
if matches!(gate, GateResult::Require) {
let path = args["path"].as_str().unwrap_or("");
let guard = self.run_context.read().await;
let dflt = RunContext::default();
let rc = guard.as_ref().unwrap_or(&dflt);
let pre_allowed = if is_file_read_tool(tool_name) {
rc.is_read_allowed(path)
} else if is_file_write_tool(tool_name) {
rc.is_write_allowed(path)
} else {
false
};
if pre_allowed { gate = GateResult::Allow; }
}
match gate {
GateResult::Allow => Ok(GateOutcome::Proceed),
GateResult::Deny => {
let msg = "Tool call denied by approval policy.".to_string();
info!(session_id = self.session_id, tool = %tool_name, tool_call_id, "approval: denied");
chat_llm_tools::reject(pool, tool_call_id, &msg).await?;
em.tool_rejected(tool_call_id, msg).await;
Ok(GateOutcome::Rejected)
}
GateResult::Require => {
if self.auto_deny_approvals.load(Ordering::Relaxed) {
let msg = "Tool call auto-denied: this session does not support approval requests.".to_string();
info!(session_id = self.session_id, tool = %tool_name, tool_call_id, "auto_deny_approvals: denied");
chat_llm_tools::reject(pool, tool_call_id, &msg).await?;
em.tool_rejected(tool_call_id, msg).await;
return Ok(GateOutcome::Rejected);
}
// Mark as pending before suspending so restart/refresh shows the
// approval form (not "Interrupted") and auto-resume re-gates.
chat_llm_tools::set_approval_pending(pool, tool_call_id).await?;
let ctx_label = self.context_label.read().ok().and_then(|g| g.clone());
let (request_id, approve_rx) = self.approval.register(
self.session_id, tool_call_id, tool_name,
args.clone(), agent_id, &self.source,
ctx_label.as_deref(), category,
).await;
info!(session_id = self.session_id, tool = %tool_name, tool_call_id, request_id, "approval: waiting for human");
self.emit_approval_event(em, request_id, tool_call_id, tool_name, args).await;
match approve_rx.await {
Ok(ApprovalDecision::Approved) => {
info!(session_id = self.session_id, request_id, tool = %tool_name, "approval: approved");
Ok(GateOutcome::Proceed)
}
Ok(ApprovalDecision::Rejected { note }) => {
info!(session_id = self.session_id, request_id, tool = %tool_name, %note, "approval: rejected");
let msg = ApprovalDecision::rejection_message(&note);
chat_llm_tools::reject(pool, tool_call_id, &msg).await?;
em.tool_rejected(tool_call_id, msg).await;
Ok(GateOutcome::Rejected)
}
Err(_) => {
// WS closed while waiting — session is orphaned.
warn!(session_id = self.session_id, request_id, "approval channel closed (WS disconnected), aborting");
Ok(GateOutcome::ChannelClosed)
}
}
}
}
}
}
@@ -0,0 +1,168 @@
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};
use serde_json::Value;
use crate::mcp::McpManager;
use crate::tools::Tool;
use crate::tools::tool_names as tn;
pub use core_api::interface_tool::{InterfaceTool, ToolFuture};
/// All configuration for a single agent run (root or sub-agent).
///
/// Passed by reference to `run_agent_turn` and `dispatch_call_agent`.
/// Callers build this once in `handle_message`; sub-agents receive a derived
/// config with an empty `interface_tools` (except `activate_tools`) and fresh
/// `active_mcp_grants`.
pub struct AgentRunConfig {
pub agent_id: String,
pub client_name: String,
/// Recursion depth: 0 = root agent, 1+ = sub-agent.
pub depth: i64,
/// Global tool definitions (built-in tools only, no MCP, **no `Config` category**).
/// MCP tools and the `Config` group are included dynamically in `all_tool_defs()`
/// based on `active_mcp_grants`.
pub base_tool_defs: Vec<Value>,
/// Definitions of the built-in `Config`-category tools (the lazy `config` group).
/// Appended by `all_tool_defs()` only when `active_mcp_grants` contains `"config"`.
/// Already filtered (interactive-only / approval visibility) by the builder.
pub config_tool_defs: Vec<Value>,
/// Static extra context injected into the first (cacheable) system message.
/// Example: Telegram HTML format instructions. Should never contain
/// per-turn data (timestamps, user-specific state) so the cached prefix
/// remains byte-identical across turns.
pub extra_system: Option<String>,
/// Dynamic extra context injected as a separate system message AFTER the
/// conversation history, just before the LLM generates its response.
/// Example: Honcho long-term memory retrieved fresh every turn.
/// Placing it at the tail keeps the stable prefix maximally cacheable
/// while giving the model fresh user context at generation time.
pub extra_system_dynamic: Option<String>,
/// Short reminder injected as a trailing `system` message in the message list.
pub tail_reminder: Option<String>,
/// Named substitutions applied to the agent's system prompt at build time.
/// Each entry replaces `__KEY__` sentinels produced by `agents::resolve_includes`.
pub system_substitutions: HashMap<String, String>,
/// Interface-specific tools.
/// For sub-agents this contains only `activate_tools`; all others are dropped.
pub interface_tools: Vec<InterfaceTool>,
/// Tools provided by the active memory backend (e.g. `memory_query`).
pub memory_tools: Vec<Arc<dyn Tool>>,
/// Image generation tools — present only when at least one provider is registered.
pub image_tools: Vec<Arc<dyn Tool>>,
/// MCP manager — used by `all_tool_defs()` to resolve which tools to include.
pub mcp: Arc<McpManager>,
/// Set of MCP server names currently granted (activated) for this agent run.
///
/// - Root agents: pre-populated from `session_mcp_grants` DB at config-build time;
/// updated in-place by `activate_tools`.
/// - Sub-agents: starts empty; populated by `activate_tools` (stack-scoped, no
/// session leak); deleted from DB when the stack frame terminates.
///
/// May also contain the reserved keyword `"config"`, which unlocks the built-in
/// `Config`-category tools (`config_tool_defs`) rather than an MCP server.
///
/// `all_tool_defs()` re-reads this set on every call, so tools activated via
/// `activate_tools` in round N are available in round N+1 within the same turn.
pub active_mcp_grants: Arc<RwLock<HashSet<String>>>,
/// Tool names that are restricted to the root agent (depth == 0).
/// Filtered out when deriving a sub-agent config via `for_sub_agent()`.
pub root_only_tool_names: Vec<String>,
}
impl AgentRunConfig {
/// Full tool list sent to the LLM on each round:
/// base tools + MCP tools for granted servers (dynamic) + `config` group (if granted)
/// + memory tools + interface tools.
///
/// Dynamic groups are re-queried every call so that an `activate_tools` call in
/// round N makes the tools visible in round N+1 without rebuilding the whole config.
pub fn all_tool_defs(&self) -> Vec<Value> {
let mut defs = self.base_tool_defs.clone();
// Dynamic groups: read the currently-granted set (MCP server names + `config`).
let granted: HashSet<String> = self.active_mcp_grants
.read()
.map(|g| g.clone())
.unwrap_or_default();
// MCP servers: include tools for the granted server names.
let servers: Vec<String> = granted.iter()
.filter(|n| n.as_str() != crate::tools::tool_names::CONFIG_GROUP)
.cloned()
.collect();
if !servers.is_empty() {
defs.extend(
self.mcp.tools_for(&servers)
.iter()
.map(|t| t.to_openai_definition()),
);
}
// `config` group: include the built-in Config-category tools on demand.
if granted.contains(crate::tools::tool_names::CONFIG_GROUP) {
defs.extend(self.config_tool_defs.iter().cloned());
}
defs.extend(self.memory_tools.iter().map(|t| t.openai_definition()));
defs.extend(self.image_tools.iter().map(|t| t.openai_definition()));
defs.extend(self.interface_tools.iter().map(|t| t.definition.clone()));
defs
}
/// Derives a config for a sub-agent:
/// - Inherits base tools, memory tools, and MCP manager.
/// - Starts with **empty** `active_mcp_grants` (sub-agents activate what they need).
/// - Drops all interface tools (caller re-injects `activate_tools` explicitly).
/// - Increments depth.
pub fn for_sub_agent(&self, agent_id: String, client_name: String) -> Self {
let root_only = |defs: &mut Vec<Value>| {
defs.retain(|def| {
let name = def["function"]["name"].as_str().unwrap_or("");
!self.root_only_tool_names.iter().any(|n| n == name)
});
};
let mut defs = self.base_tool_defs.clone();
root_only(&mut defs);
// Strip the per-level augmentations that the config builders re-derive, so
// they are never inherited: `ask_user_clarification` is added by
// `build_agent_config` (root) and re-added by `dispatch_sub_agent`;
// `execute_subtask` is added by `dispatch_sub_agent`. Leaving them in the
// inherited set would duplicate them (depth ≥ 1 for `ask_user_clarification`,
// depth ≥ 2 for `execute_subtask`) and the OpenAI-compat APIs reject
// non-unique tool names with HTTP 400. With this strip, `dispatch_sub_agent`
// is the single owner of sub-agent augmentation and duplication is
// structurally impossible — no dedup pass needed anywhere.
{
const RE_DERIVED: &[&str] = &[tn::ASK_USER_CLARIFICATION, tn::EXECUTE_SUBTASK];
defs.retain(|d| {
let name = d["function"]["name"].as_str().unwrap_or("");
!RE_DERIVED.contains(&name)
});
}
// Inherit the (already filtered) `config` group, dropping any root-only tool.
let mut config_defs = self.config_tool_defs.clone();
root_only(&mut config_defs);
Self {
agent_id,
client_name,
depth: self.depth + 1,
base_tool_defs: defs,
config_tool_defs: config_defs,
extra_system: None,
extra_system_dynamic: None,
tail_reminder: None,
system_substitutions: HashMap::new(),
interface_tools: vec![],
memory_tools: self.memory_tools.clone(),
image_tools: self.image_tools.clone(),
mcp: Arc::clone(&self.mcp),
active_mcp_grants: Arc::new(RwLock::new(HashSet::new())),
root_only_tool_names: self.root_only_tool_names.clone(),
}
}
}
@@ -0,0 +1,138 @@
//! One LLM call per round, with automatic model fallback.
//!
//! Extracted from `run_agent_turn`: on a retriable error (5xx / network) it retries
//! up to `MAX_LLM_ATTEMPTS` models in priority order, rebuilding the message list
//! when the replacement model has a different `prompt_cache` setting, and emits
//! `ModelFallback` / `LlmFailed` along the way.
use std::collections::HashSet;
use std::sync::Arc;
use serde_json::Value;
use tokio_util::sync::CancellationToken;
use tracing::{error, warn};
use crate::chatbot::{ChatOptions, LlmTurn};
use crate::llm::{LlmEntry, LlmStrength};
use super::ChatSessionHandler;
use super::emitter::TurnEmitter;
use super::interface_tools::AgentRunConfig;
/// Outcome of one round's LLM call.
pub(super) enum RoundLlm {
/// The model responded (message or tool calls).
Turn(LlmTurn),
/// The turn was cancelled (`/stop`) while the request was in flight.
Cancelled,
/// All fallback attempts were exhausted, or an error is non-retriable.
Failed(anyhow::Error),
}
/// Maximum number of models tried in one round before giving up.
const MAX_LLM_ATTEMPTS: usize = 3;
impl ChatSessionHandler {
/// Calls the current model and, on a retriable failure, falls back to the next
/// model in priority order. Mutates `cur_name` / `cur_llm` / `messages` in place
/// so the caller keeps using the model that actually produced the turn.
#[allow(clippy::too_many_arguments)]
pub(super) async fn call_llm_round(
&self,
stack_id: i64,
config: &AgentRunConfig,
active_grants: &HashSet<String>,
tool_defs: &[Value],
req_scope: Option<&str>,
req_strength: Option<LlmStrength>,
cur_name: &mut String,
cur_llm: &mut Arc<LlmEntry>,
messages: &mut Vec<Value>,
token: &CancellationToken,
em: &TurnEmitter<'_>,
) -> RoundLlm {
let mut tried_this_round: Vec<String> = vec![cur_name.clone()];
loop {
let options = ChatOptions {
model: cur_llm.model.clone(),
max_tokens: None,
temperature: None,
session_id: Some(self.session_id),
stack_id: Some(stack_id),
};
// Clone the Arc so the in-flight future does not borrow `cur_llm` across
// the fallback reassignment below. On cancel we drop the future
// (aborting the request) and return immediately.
let client = cur_llm.client.clone();
let call_result = tokio::select! {
_ = token.cancelled() => return RoundLlm::Cancelled,
r = client.chat_with_tools(messages.as_slice(), tool_defs, &options) => r,
};
let e = match call_result {
Ok(t) => {
self.llm_manager.mark_success(cur_name).await;
return RoundLlm::Turn(t);
}
Err(e) => e,
};
error!(session_id = self.session_id, client = %cur_name, error = %e, "LLM call failed");
self.llm_manager.mark_failure(cur_name, &e.to_string()).await;
let can_fallback = tried_this_round.len() < MAX_LLM_ATTEMPTS
&& is_retriable_llm_error(&e);
if !can_fallback {
em.llm_failed(tried_this_round.clone(), e.to_string()).await;
return RoundLlm::Failed(e);
}
let excluded: Vec<&str> = tried_this_round.iter().map(String::as_str).collect();
match self.llm_manager.select_excluding(&excluded, req_scope, req_strength).await {
Ok((next_name, next_llm)) => {
warn!(session_id = self.session_id, from = %cur_name, to = %next_name, "LLM fallback");
em.model_fallback(cur_name.clone(), next_name.clone(), first_line(&e.to_string())).await;
tried_this_round.push(next_name.clone());
*cur_name = next_name;
*cur_llm = next_llm;
// Rebuild messages if the new model uses different prompt_cache
// settings (e.g. switching from OpenRouter/Anthropic to DeepSeek).
match self.build_openai_messages(
&self.db, stack_id, &config.agent_id,
config.extra_system.as_deref(), config.extra_system_dynamic.as_deref(),
config.tail_reminder.as_deref(), active_grants,
&config.system_substitutions, cur_llm.prompt_cache,
).await {
Ok(m) => *messages = m,
Err(e) => return RoundLlm::Failed(e),
}
}
Err(_) => {
em.llm_failed(tried_this_round.clone(), e.to_string()).await;
return RoundLlm::Failed(e);
}
}
}
}
}
/// Whether an LLM error is worth retrying on a different model.
fn is_retriable_llm_error(e: &anyhow::Error) -> bool {
let msg = e.to_string().to_lowercase();
// Never retry client errors — the request itself is malformed or unauthorized.
// 400 is excluded: some providers reject valid requests that others accept
// (e.g. DeepSeek requires reasoning_content echo, OpenAI does not), so
// retrying on a different model can succeed.
for code in ["401", "403", "404", "422"] {
if msg.contains(code) {
return false;
}
}
true
}
fn first_line(s: &str) -> String {
s.lines().next().unwrap_or(s).to_string()
}
@@ -0,0 +1,431 @@
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, trace};
use crate::tools::tool_names as tn;
use crate::chat_event_bus::ToolCallEvent;
use crate::chatbot::{LlmTurn, ToolCall};
use crate::db::{chat_history, chat_llm_tools};
use crate::events::ServerEvent;
use crate::tools::{
ExecutionOutcome, SimpleExecution, ToolDescriptionLength, ToolExecution, ToolResult,
};
use futures::stream::{self, StreamExt};
use super::{ChatSessionHandler, PendingUserInput, TurnOutcome};
use super::dispatch::{is_sync_sub_agent, DispatchResult};
use super::emitter::TurnEmitter;
use super::gate::GateOutcome;
use super::llm_call::RoundLlm;
use super::outcome::RecordFlow;
use super::interface_tools::AgentRunConfig;
/// Whether, after handling one tool call, the round loop should continue to the
/// next call or the whole turn should end.
enum CallFlow {
Continue,
End(TurnOutcome),
}
/// Outcome of gating + dispatching one call inside a concurrent sub-agent batch,
/// carried from the concurrent phase to the ordered recording phase.
enum GatedExec {
/// Gate passed; the sub-agent produced an outcome to record. `effective` is the
/// working-dir-resolved args used for recording (FileChanged / logging).
Done { effective: serde_json::Value, outcome: ExecutionOutcome },
/// Approval gate rejected the call — already marked/emitted by the gate; skip it.
Rejected,
/// The turn must end now: the clarification WS channel closed (dispatch returned
/// `AbortPending`) or the approval gate's channel closed.
AbortTurn,
}
impl ChatSessionHandler {
/// Inner loop of an agent (root or sub). Persists messages to `stack_id`,
/// emits Thinking/ToolStart/ToolDone/PendingWrite/ApprovalRequired/AgentStart/AgentDone events.
/// Returns the outcome; the caller decides what to emit on completion
/// (Done for root, AgentDone+tool-result for sub-agents).
pub(super) fn run_agent_turn<'a>(
&'a self,
stack_id: i64,
config: &'a AgentRunConfig,
token: &'a CancellationToken,
tx: &'a mpsc::Sender<ServerEvent>,
// Queued user input for live injection (root interactive turn only).
// `None` for sub-agents / resume / non-interactive runners.
pending_input: Option<&'a Arc<dyn PendingUserInput>>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = anyhow::Result<TurnOutcome>> + Send + 'a>> {
Box::pin(async move {
let pool = &self.db;
let em = TurnEmitter::new(tx);
// Resolve the initial model. `cur_name`/`cur_llm` are updated in-place
// when the fallback logic switches to a different model mid-turn.
let mut cur_name = config.client_name.clone();
let mut cur_llm = self.llm_manager.get(&cur_name).await
.ok_or_else(|| anyhow::anyhow!("LLM client '{}' not found", cur_name))?;
// Scope/strength needed for fallback re-selection.
let meta = crate::agents::load_meta(&config.agent_id).ok();
let req_scope = meta.as_ref().and_then(|m| m.scope.as_deref()).map(str::to_string);
let req_strength = meta.as_ref().and_then(|m| m.strength);
// Accumulates tool calls across all rounds for the event bus.
let mut all_tool_calls: Vec<ToolCallEvent> = Vec::new();
for round in 0..self.max_tool_rounds {
if token.is_cancelled() {
return Ok(TurnOutcome::Cancelled);
}
// ── Live user-message injection ─────────────────────────────────────
// A round boundary is the one clean ordering point: the previous
// round's assistant message + tool results are all persisted, so a
// `user` row appended here is well-ordered. Each queued message is
// saved individually and echoed (telnet-style: the bubble appears only
// now), then picked up by `build_openai_messages` below in this same
// round — so the model sees it immediately. The MessageBuilder merges
// consecutive user rows into one `role:user` for the LLM. Does not
// reset the round budget. Only ever `Some` for the root interactive turn.
if let Some(input) = pending_input {
for msg in input.drain_user().await {
let attachments = msg.metadata.as_ref()
.map(|m| m.attachments.clone())
.unwrap_or_default();
// A custom slash command persists its expanded template (for LLM
// replay) but the bubble must show the typed command — emit the
// command's `display` form when present.
let echo = msg.metadata.as_ref()
.and_then(|m| m.command.as_ref())
.map(|c| c.display.clone())
.unwrap_or_else(|| msg.content.clone());
let id = chat_history::append_with_metadata(
pool, stack_id, &chat_history::Role::User,
&msg.content, false, None, msg.metadata.as_ref(),
).await?;
em.user_message(id, echo, attachments).await;
}
}
trace!(session_id = self.session_id, stack_id, agent_id = config.agent_id, round, "starting round");
let active_grants_snapshot = config.active_mcp_grants
.read()
.map(|g| g.clone())
.unwrap_or_default();
// Messages are (re)built with the current model's prompt_cache flag.
// On fallback within the same round `call_llm_round` rebuilds them again
// if the replacement model has a different prompt_cache setting.
let mut messages = self.build_openai_messages(pool, stack_id, &config.agent_id, config.extra_system.as_deref(), config.extra_system_dynamic.as_deref(), config.tail_reminder.as_deref(), &active_grants_snapshot, &config.system_substitutions, cur_llm.prompt_cache).await?;
let tool_defs = config.all_tool_defs();
// Record every tool actually offered to the LLM so the Security-groups
// UI can list/gate dynamically-injected tools. Cheap no-op once each
// name is known; new names are persisted off the turn's critical path.
self.tool_discovery.observe(&tool_defs);
// One LLM call for this round, with automatic model fallback on
// retriable errors. `cur_name`/`cur_llm`/`messages` are updated in place.
let turn_result = match self.call_llm_round(
stack_id, config, &active_grants_snapshot, &tool_defs,
req_scope.as_deref(), req_strength,
&mut cur_name, &mut cur_llm, &mut messages, token, &em,
).await {
RoundLlm::Turn(t) => t,
RoundLlm::Cancelled => return Ok(TurnOutcome::Cancelled),
RoundLlm::Failed(e) => return Err(e),
};
match turn_result {
LlmTurn::Message(resp) => {
let message_id = chat_history::append(
pool, stack_id, &chat_history::Role::Assistant, &resp.content, false,
resp.reasoning_content.as_deref(),
).await?;
if let (Some(i), Some(o)) = (resp.input_tokens, resp.output_tokens) {
chat_history::set_usage(pool, message_id, i, o, 0, resp.cost).await?;
}
return Ok(TurnOutcome::Final {
content: resp.content,
message_id,
input_tokens: resp.input_tokens,
output_tokens: resp.output_tokens,
truncated: resp.truncated,
tool_calls: all_tool_calls,
});
}
LlmTurn::ToolCalls { content: assistant_text, calls, input_tokens, output_tokens, reasoning_content, cost, .. } => {
let message_id = chat_history::append(
pool, stack_id, &chat_history::Role::Assistant, &assistant_text, false,
reasoning_content.as_deref(),
).await?;
if let (Some(i), Some(o)) = (input_tokens, output_tokens) {
chat_history::set_usage(pool, message_id, i, o, 0, cost).await?;
}
if !assistant_text.trim().is_empty() || input_tokens.is_some() {
em.thinking(message_id, assistant_text, input_tokens, output_tokens).await;
}
// A homogeneous batch of ≥2 synchronous sub-agent calls is fanned
// out concurrently (bounded by `max_parallel_subagents`). Any other
// shape — a single call, or a mix with regular tools — keeps the
// strictly sequential path, so tool ordering and side-effects are
// unchanged for everything except this well-defined case.
if calls.len() >= 2 && calls.iter().all(|c| is_sync_sub_agent(&c.name, &c.arguments)) {
match self.handle_sub_agent_batch(
stack_id, config, message_id, &calls, token, tx, &em, &mut all_tool_calls,
).await? {
CallFlow::Continue => {}
CallFlow::End(outcome) => return Ok(outcome),
}
} else {
for call in &calls {
// Stop before each call so a /stop (or a cancelled sub-agent,
// which shares this token) aborts the rest of the round.
if token.is_cancelled() {
return Ok(TurnOutcome::Cancelled);
}
match self.handle_tool_call(
stack_id, config, message_id, call, token, tx, &em, &mut all_tool_calls,
).await? {
CallFlow::Continue => {}
CallFlow::End(outcome) => return Ok(outcome),
}
}
}
}
}
}
Ok(TurnOutcome::Exhausted)
}) // end Box::pin
}
/// Handles a single tool call within a round: persists the call row, emits
/// `ToolStart`, resolves the working directory, runs the approval gate, handles
/// `restart`, dispatches, and records the outcome. Returns [`CallFlow::Continue`]
/// to move on to the next call, or [`CallFlow::End`] to end the whole turn.
#[allow(clippy::too_many_arguments)]
async fn handle_tool_call(
&self,
stack_id: i64,
config: &AgentRunConfig,
message_id: i64,
call: &ToolCall,
token: &CancellationToken,
tx: &mpsc::Sender<ServerEvent>,
em: &TurnEmitter<'_>,
all_tool_calls: &mut Vec<ToolCallEvent>,
) -> anyhow::Result<CallFlow> {
let pool = &self.db;
let args_str = serde_json::to_string(&call.arguments)
.unwrap_or_else(|_| "{}".to_string());
let tool_call_id = chat_llm_tools::append(pool, message_id, &call.name, &args_str).await?;
em.tool_start(
tool_call_id, message_id,
call.name.clone(),
call.arguments.clone(),
self.tools.describe_call(&call.name, &call.arguments, ToolDescriptionLength::Short),
self.tools.describe_call(&call.name, &call.arguments, ToolDescriptionLength::Full),
self.tools.target_path(&call.name, &call.arguments),
).await;
// Resolve relative paths / inject workdir from the RunContext.
// `call.arguments` (originals) were used for the ToolStart event and DB
// logging above; `effective_args` is used from here on.
let effective_args = self.effective_args(&call.name, &call.arguments).await;
match self.run_approval_gate(tool_call_id, &call.name, &effective_args, &config.agent_id, em).await? {
GateOutcome::Proceed => {}
GateOutcome::Rejected => return Ok(CallFlow::Continue),
GateOutcome::ChannelClosed => return Ok(CallFlow::End(TurnOutcome::Cancelled)),
}
debug!(session_id = self.session_id, tool = %call.name, tool_call_id, "dispatching");
// `restart` calls process::exit — mark the call done in the DB first so it
// doesn't reappear as `pending` after the supervisor relaunches.
if call.name == tn::RESTART {
info!(session_id = self.session_id, tool_call_id, "restart approved — marking done then exiting");
chat_llm_tools::complete(pool, tool_call_id, "Riavvio avviato.", "string").await?;
em.tool_done(tool_call_id, "Riavvio avviato.".to_string(), "string".to_string()).await;
// Use _exit() to skip C atexit handlers (e.g. Metal GPU cleanup in
// whisper-rs/ggml, which aborts with SIGABRT and yields exit code 134
// instead of 255 — breaking the run.sh restart supervisor).
unsafe { libc::_exit(-1) }
}
// Route the approved call to its executor. `AbortPending` means the
// clarification WS channel closed — end the turn and leave the tool
// `pending` for resume to re-ask.
let outcome = match self.execute_tool_call(
stack_id, config, tool_call_id, &call.name, &effective_args, token, tx,
).await {
DispatchResult::Outcome(o) => o,
DispatchResult::AbortPending => return Ok(CallFlow::End(TurnOutcome::Cancelled)),
};
match self.record_tool_outcome(
tool_call_id, &call.name, &effective_args, outcome, em, Some(all_tool_calls),
).await? {
RecordFlow::Continue => Ok(CallFlow::Continue),
RecordFlow::Abort => Ok(CallFlow::End(TurnOutcome::Cancelled)),
}
}
/// Concurrent variant of the tool-call loop for a homogeneous batch of
/// synchronous sub-agent calls (`execute_task` mode=sync / `execute_subtask`).
/// Only called when every call in the round is such a sub-agent (see the
/// dispatch in `run_agent_turn`), so `restart` and side-effecting tools can
/// never appear here and the sequential path is left byte-for-byte intact.
///
/// Ordering invariant: the LLM reconstructs tool results by autoincrement id
/// (`chat_llm_tools ORDER BY id ASC`). **Phase 1** therefore allocates every
/// call's row in `calls` order *before* any concurrent work, so completion
/// order is irrelevant. **Phase 2** runs the approval gate + dispatch for all
/// calls concurrently, bounded by `max_parallel_subagents`. **Phase 3** records
/// the outcomes back in `calls` order, so `all_tool_calls` ordering and the
/// shared-token cancellation semantics match the sequential path.
#[allow(clippy::too_many_arguments)]
async fn handle_sub_agent_batch(
&self,
stack_id: i64,
config: &AgentRunConfig,
message_id: i64,
calls: &[ToolCall],
token: &CancellationToken,
tx: &mpsc::Sender<ServerEvent>,
em: &TurnEmitter<'_>,
all_tool_calls: &mut Vec<ToolCallEvent>,
) -> anyhow::Result<CallFlow> {
let pool = &self.db;
// ── Phase 1: allocate tool_call_id rows in `calls` order ────────────────────
// The id fixes the LLM-visible order regardless of which sub-agent finishes
// first, so this pre-pass MUST stay sequential and precede the fan-out.
let mut started: Vec<(&ToolCall, i64)> = Vec::with_capacity(calls.len());
for call in calls {
let args_str = serde_json::to_string(&call.arguments)
.unwrap_or_else(|_| "{}".to_string());
let tool_call_id = chat_llm_tools::append(pool, message_id, &call.name, &args_str).await?;
em.tool_start(
tool_call_id, message_id,
call.name.clone(),
call.arguments.clone(),
self.tools.describe_call(&call.name, &call.arguments, ToolDescriptionLength::Short),
self.tools.describe_call(&call.name, &call.arguments, ToolDescriptionLength::Full),
self.tools.target_path(&call.name, &call.arguments),
).await;
started.push((call, tool_call_id));
}
// ── Phase 2: gate + dispatch concurrently, bounded ──────────────────────────
// Every future borrows `&self`/`config`/`token`/`tx`/`em` (all shared refs)
// and writes only to its own distinct child stack + tool_call_id, so there is
// no shared mutable state between siblings. Results are keyed back by index.
let limit = self.max_parallel_subagents.max(1);
let mut results: Vec<Option<GatedExec>> = (0..started.len()).map(|_| None).collect();
// Feed the stream fully-owned items `(idx, tool_call_id, name, arguments)`.
// Passing a borrowed `&ToolCall` as the closure input makes the returned async
// block's lifetime higher-ranked ("FnOnce is not general enough"); owning the
// per-call data means each future only borrows `self`/`config`/`token`/`tx`/`em`
// from the enclosing scope, all at the single concrete turn lifetime.
let jobs: Vec<(usize, i64, String, serde_json::Value)> = started.iter().enumerate()
.map(|(idx, (call, id))| (idx, *id, call.name.clone(), call.arguments.clone()))
.collect();
{
let mut stream = stream::iter(jobs)
.map(|(idx, tool_call_id, name, arguments)| async move {
let effective = self.effective_args(&name, &arguments).await;
let gated = match self.run_approval_gate(
tool_call_id, &name, &effective, &config.agent_id, em,
).await {
Ok(GateOutcome::Proceed) => match self.execute_tool_call(
stack_id, config, tool_call_id, &name, &effective, token, tx,
).await {
DispatchResult::Outcome(outcome) => Ok(GatedExec::Done { effective, outcome }),
DispatchResult::AbortPending => Ok(GatedExec::AbortTurn),
},
Ok(GateOutcome::Rejected) => Ok(GatedExec::Rejected),
Ok(GateOutcome::ChannelClosed) => Ok(GatedExec::AbortTurn),
Err(e) => Err(e),
};
(idx, gated)
})
.buffer_unordered(limit);
while let Some((idx, gated)) = stream.next().await {
results[idx] = Some(gated?);
}
}
// ── Phase 3: record outcomes in `calls` order ───────────────────────────────
let mut abort = false;
for (idx, (call, tool_call_id)) in started.iter().enumerate() {
match results[idx].take().expect("every started sub-agent call produced a result") {
// The gate already marked the row rejected and emitted the event.
GatedExec::Rejected => {}
GatedExec::AbortTurn => abort = true,
GatedExec::Done { effective, outcome } => {
match self.record_tool_outcome(
*tool_call_id, &call.name, &effective, outcome, em, Some(all_tool_calls),
).await? {
RecordFlow::Continue => {}
RecordFlow::Abort => abort = true,
}
}
}
}
// The shared token means a /stop (or a cancelled sibling) has already stopped
// the others; ending the turn here mirrors the sequential path's early return.
if abort || token.is_cancelled() {
Ok(CallFlow::End(TurnOutcome::Cancelled))
} else {
Ok(CallFlow::Continue)
}
}
/// Builds a [`ToolExecution`] for a single tool call, covering every tool that
/// flows through the unified (cancellable) dispatch path: interface tools,
/// memory/image tools, MCP tools, and the built-in registry (incl.
/// `execute_cmd`). Returns `None` only for an unknown tool name. The handle
/// borrows `self` and `config`, both of which outlive the turn.
pub(super) fn build_execution<'a>(
&'a self,
name: &str,
args: serde_json::Value,
config: &'a AgentRunConfig,
) -> Option<Box<dyn ToolExecution + 'a>> {
// Interface tools (closures injected per-interface, e.g. activate_tools).
if let Some(tool) = config.interface_tools.iter().find(|t| t.name() == name) {
let handler = std::sync::Arc::clone(&tool.handler);
return Some(Box::new(SimpleExecution::new(
Box::pin(async move { handler(args).await.map(ToolResult::Text) }),
)));
}
// Memory + image tools (registered ad-hoc on the config).
if let Some(tool) = config.memory_tools.iter().find(|t| t.name() == name) {
return Some(tool.run(args));
}
if let Some(tool) = config.image_tools.iter().find(|t| t.name() == name) {
return Some(tool.run(args));
}
// MCP tools (`server::tool`). Clone the Arc so the work future is 'static.
if let Some((srv, mcp_tool)) = crate::mcp::parse_mcp_tool_name(name) {
let mcp = std::sync::Arc::clone(&self.mcp);
let srv = srv.to_string();
let mcp_tool = mcp_tool.to_string();
let fut: std::pin::Pin<Box<dyn std::future::Future<Output = anyhow::Result<ToolResult>> + Send>> =
Box::pin(async move { mcp.call(&srv, &mcp_tool, args).await });
return Some(Box::new(SimpleExecution::new(fut)));
}
// Built-in registry tools (incl. execute_cmd, whose SimpleExecution kills
// the child via kill_on_drop when the work future is dropped on /stop).
self.tools.run(name, args)
}
}
@@ -0,0 +1,567 @@
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use serde_json::{Value, json};
use sqlx::SqlitePool;
use crate::compactor::{ContextCompactor, SUMMARY_PREFIX};
use crate::config::DatetimeConfig;
use crate::db::{chat_history, chat_llm_tools, chat_summaries};
use crate::mcp::McpManager;
use crate::tools::tool_names as tn;
/// Registry of installed skills, relative to Skald's process cwd. Injected into agents
/// that have `inject_skills` enabled (the default).
const SKILLS_INDEX_PATH: &str = "skills/index.md";
/// OS description (type + version), computed once — it does not change at runtime.
fn os_description() -> &'static str {
static OS: std::sync::OnceLock<String> = std::sync::OnceLock::new();
OS.get_or_init(|| os_info::get().to_string())
}
/// System IANA timezone name (e.g. `Europe/Rome`), computed once. `None` if it can't
/// be determined.
fn system_timezone() -> Option<&'static str> {
static TZ: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
TZ.get_or_init(|| iana_time_zone::get_timezone().ok()).as_deref()
}
/// Pure service that builds the OpenAI-format message array for one LLM round.
///
/// Extracting this from `ChatSessionHandler` allows the builder to be constructed
/// and called in isolation (e.g. in integration tests with an in-memory SQLite DB)
/// without needing the full handler and all its dependencies.
pub struct MessageBuilder {
pub pool: Arc<SqlitePool>,
pub session_id: i64,
pub mcp: Arc<McpManager>,
pub datetime_config: DatetimeConfig,
pub max_history_messages: usize,
pub max_tool_result_chars: Option<usize>,
pub compactor: Option<Arc<ContextCompactor>>,
/// Effective working directory for this session. When set (e.g. from a project
/// RunContext), it overrides the process cwd in the date/time/OS/WD tail block.
pub working_directory: Option<std::path::PathBuf>,
}
impl MessageBuilder {
/// Builds a raw OpenAI-format message array from the persisted history,
/// reconstructing assistant tool-call entries and tool-result entries from
/// the `chat_llm_tools` table.
///
/// `active_mcp_grants` is the set of MCP server names currently granted for
/// this session. It is used to build the compact MCP availability list injected
/// into the system prompt so the LLM knows which servers it can activate.
///
/// ## Message order (optimised for prefix KV caching)
///
/// ```text
/// 1. [system] Static content — AGENT.md + memory files + extra_system_static + MCP list
/// Tagged cache_control:ephemeral when cache_hints=true (Anthropic via OpenRouter).
///
/// 2. [system] Scratchpad — emitted only when non-empty, BEFORE the conversation.
///
/// 3. [system] Compaction summary — if a summary exists for this stack.
///
/// 4. [user / assistant / tool] Conversation history.
///
/// 5. [system] Dynamic tail — extra_system_dynamic + current date/time/OS/cwd.
///
/// 6. [system] Tail reminder — short anti-drift reminder (e.g. Telegram format).
/// ```
pub async fn build(
&self,
stack_id: i64,
agent_id: &str,
extra_system_static: Option<&str>,
extra_system_dynamic: Option<&str>,
tail_reminder: Option<&str>,
active_mcp_grants: &HashSet<String>,
system_substitutions: &HashMap<String, String>,
cache_hints: bool,
) -> anyhow::Result<Vec<Value>> {
let pool = &*self.pool;
// ── 1. Static system message ──────────────────────────────────────────
let mut static_content = crate::agents::load_prompt(agent_id)?;
let meta = crate::agents::load_meta(agent_id)?;
if !meta.inject_memory.is_empty() {
static_content.push_str(
"\n\n---\nThe following memory files have been loaded automatically. \
You can edit them with `edit_file` or `write_file` using the path shown.\n"
);
for mem_path in &meta.inject_memory {
// Resolve the entry to (absolute path to read, path to show the agent).
let (abs, display) = self.resolve_memory_path(mem_path);
let content = tokio::fs::read_to_string(&abs).await.ok();
match content {
Some(c) => static_content.push_str(&format!(
"\n<memory_file path=\"{display}\">\n{c}\n</memory_file>\n"
)),
None => static_content.push_str(&format!(
"\n<memory_file path=\"{display}\">\n(file not created yet)\n</memory_file>\n"
)),
}
}
}
// ── Skills index ──────────────────────────────────────────────────────
// Injected for every agent unless it opts out (`inject_skills: false`).
// Reuses the memory-path resolution so the shown path is relative when the
// index is under the session WD, absolute otherwise (it lives under Skald's
// own cwd, so it shows as absolute inside project sessions). Skipped silently
// when no skills are installed.
if meta.inject_skills {
let (abs, display) = self.resolve_memory_path(SKILLS_INDEX_PATH);
if let Ok(c) = tokio::fs::read_to_string(&abs).await {
static_content.push_str(&format!(
"\n\n---\nInstalled skills you can use (read the linked `SKILL.md` before running a skill):\n\
\n<skills_index path=\"{display}\">\n{c}\n</skills_index>\n"
));
}
}
if let Some(extra) = extra_system_static {
static_content.push_str("\n\n---\n");
static_content.push_str(extra);
}
if static_content.contains("__MCP_LIST__") {
static_content = static_content.replace(
"__MCP_LIST__",
&self.render_mcp_list(active_mcp_grants),
);
}
for (key, value) in system_substitutions {
let sentinel = format!("__{key}__");
if static_content.contains(sentinel.as_str()) {
static_content = static_content.replace(sentinel.as_str(), value);
}
}
let static_msg = if cache_hints {
json!({
"role": "system",
"content": [{ "type": "text", "text": static_content, "cache_control": { "type": "ephemeral" } }]
})
} else {
json!({ "role": "system", "content": static_content })
};
let mut out = vec![static_msg];
// ── 2. Scratchpad system message (before conversation) ────────────────
let scratch = crate::db::scratchpad::for_session(pool, self.session_id).await?;
if !scratch.is_empty() {
let mut s = String::from(
"<scratchpad>\n \
<!-- Temporary notes shared by all agents in this session. Not persisted across sessions. -->\n"
);
for (k, v) in &scratch {
s.push_str(&format!(" <note key=\"{k}\">{v}</note>\n"));
}
s.push_str("</scratchpad>");
out.push(json!({ "role": "system", "content": s }));
}
// ── 3. Context compaction: inject summary + load messages after boundary ──
let summary = chat_summaries::latest_for_stack(pool, stack_id).await?;
let mut history = match &summary {
Some(s) => {
out.push(json!({
"role": "system",
"content": format!(
"{SUMMARY_PREFIX}\n\n{}\n\n\
[End of context summary — the following messages are the most recent exchanges in full.]",
s.content
)
}));
chat_history::for_stack_since(pool, stack_id, s.covers_up_to_message_id).await?
}
None => chat_history::for_stack(pool, stack_id).await?,
};
if self.compactor.is_none() && history.len() > self.max_history_messages {
history.drain(..history.len() - self.max_history_messages);
if matches!(history.first().map(|m| &m.role), Some(chat_history::Role::Assistant)) {
history.drain(..1);
}
}
let current_turn_boundary = history
.iter()
.rposition(|e| matches!(e.role, chat_history::Role::User | chat_history::Role::Agent));
for (idx, entry) in history.iter().enumerate() {
let is_previous_turn = current_turn_boundary.map_or(false, |b| idx < b);
match entry.role {
chat_history::Role::User | chat_history::Role::Agent => {
// Render attachments (if any) as a textual block appended to the
// user turn, generated on the fly — never persisted as content.
let content = match &entry.metadata {
Some(meta) if !meta.attachments.is_empty() => format!(
"{}{}",
entry.content,
core_api::message_meta::attachments_block(&meta.attachments),
),
_ => entry.content.clone(),
};
// Coalesce consecutive user/agent rows into a single `role:user`
// turn. The DB keeps each message as its own row (distinct bubbles,
// per-message attachments), but the model must see one clean user
// turn — e.g. when several messages were injected back-to-back at a
// round boundary, or queued together while idle. `for_stack` already
// excludes `failed` rows, so only non-failed messages merge here.
match out.last_mut() {
Some(last) if last["role"] == "user" => {
let prev = last["content"].as_str().unwrap_or("").to_string();
last["content"] = Value::String(format!("{prev}\n\n{content}"));
}
_ => out.push(json!({ "role": "user", "content": content })),
}
}
chat_history::Role::Assistant => {
let tool_calls = chat_llm_tools::for_message(pool, entry.id).await?;
if tool_calls.is_empty() {
let mut msg = json!({ "role": "assistant", "content": entry.content });
if let Some(rc) = &entry.reasoning_content {
// Echo under both names: DeepSeek expects "reasoning_content",
// MiniMax M3 and others expect "reasoning".
msg["reasoning_content"] = rc.clone().into();
msg["reasoning"] = rc.clone().into();
}
out.push(msg);
} else {
let tc_array: Vec<Value> = tool_calls
.iter()
.map(|tc| json!({
"id": format!("tc_{}", tc.id),
"type": "function",
"function": {
"name": tc.name,
"arguments": tc.arguments.as_deref().unwrap_or("{}"),
}
}))
.collect();
let mut msg = json!({
"role": "assistant",
"content": entry.content,
"tool_calls": tc_array,
});
if let Some(rc) = &entry.reasoning_content {
// Echo under both names: DeepSeek expects "reasoning_content",
// MiniMax M3 and others expect "reasoning".
msg["reasoning_content"] = rc.clone().into();
msg["reasoning"] = rc.clone().into();
}
out.push(msg);
for tc in &tool_calls {
let result_content = match tc.status.as_str() {
"done" => tc.result.as_deref().unwrap_or("").to_string(),
"failed" => format!(
"Error: {}",
tc.result.as_deref().unwrap_or("unknown error")
),
// A human/policy rejection or a /stop cancellation is a
// deliberate, terminal outcome — surface the saved reason
// (the user's justification) so the LLM understands the
// tool did NOT run and why, instead of retrying blindly.
"rejected" => tc.result.as_deref()
.unwrap_or("User rejected this tool call.")
.to_string(),
"cancelled" => tc.result.as_deref()
.unwrap_or("Tool call was cancelled by the user.")
.to_string(),
// 'pending'/'running' left behind by a crash or a lost
// connection: the call really was interrupted mid-flight.
_ => "Error: tool call was interrupted (connection lost before user approval). Please retry the operation.".to_string(),
};
let result_content = self.maybe_hide_tool_result(
result_content,
is_previous_turn,
&tc.name,
tc.arguments.as_deref(),
);
out.push(json!({
"role": "tool",
"tool_call_id": format!("tc_{}", tc.id),
"content": result_content,
}));
}
}
}
}
}
// ── 5. Dynamic tail system message (after conversation) ──────────────
{
let datetime_line = if self.datetime_config.enabled {
let now_utc = chrono::Utc::now();
let secs = now_utc.timestamp();
let secs = match self.datetime_config.round_minutes {
Some(m) if m > 0 => {
let bucket = (m as i64) * 60;
(secs / bucket) * bucket
}
_ => secs,
};
// Effective timezone: the one configured in config.yml if set, else the
// OS timezone. When resolvable we show the IANA name alongside the offset.
let tz = self.datetime_config.timezone.as_deref()
.and_then(|s| s.parse::<chrono_tz::Tz>().ok())
.or_else(|| system_timezone().and_then(|s| s.parse::<chrono_tz::Tz>().ok()));
let (formatted, tz_name) = match tz {
Some(tz) => {
use chrono::TimeZone as _;
let f = tz.timestamp_opt(secs, 0)
.single()
.map(|dt| dt.format("%Y-%m-%dT%H:%M:%S%:z").to_string())
.unwrap_or_else(|| chrono::Local::now().format("%Y-%m-%dT%H:%M:%S%:z").to_string());
(f, Some(tz.name().to_string()))
}
None => {
let f = chrono::DateTime::from_timestamp(secs, 0)
.map(|utc| utc.with_timezone(&chrono::Local).format("%Y-%m-%dT%H:%M:%S%:z").to_string())
.unwrap_or_else(|| chrono::Local::now().format("%Y-%m-%dT%H:%M:%S%:z").to_string());
(f, None)
}
};
let date_line = match tz_name {
Some(name) => format!("Current date and time: {formatted} ({name})"),
None => format!("Current date and time: {formatted}"),
};
let cwd = self.working_directory.clone()
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default())
.display()
.to_string();
Some(format!(
"{date_line}\nOperating system: {}\nWorking directory: {cwd}\n\
Filesystem tools and execute_cmd use this working directory for relative paths — \
no need to `cd` into it first.",
os_description()
))
} else {
None
};
let tail = match (extra_system_dynamic, datetime_line.as_deref()) {
(Some(dyn_ctx), Some(dt)) => Some(format!("{dyn_ctx}\n\n---\n{dt}")),
(Some(dyn_ctx), None) => Some(dyn_ctx.to_string()),
(None, Some(dt)) => Some(dt.to_string()),
(None, None) => None,
};
if let Some(content) = tail {
out.push(json!({ "role": "system", "content": content }));
}
}
// ── 6. Tail reminder ──────────────────────────────────────────────────
if let Some(reminder) = tail_reminder {
out.push(json!({ "role": "system", "content": reminder }));
}
Ok(out)
}
/// Returns the tool result as-is, or replaces it with an informative 1-line
/// summary when the result belongs to a previous turn and exceeds `max_tool_result_chars`.
fn maybe_hide_tool_result(
&self,
result: String,
is_previous_turn: bool,
tool_name: &str,
arguments: Option<&str>,
) -> String {
if !is_previous_turn {
return result;
}
let Some(limit) = self.max_tool_result_chars else {
return result;
};
if result.len() <= limit {
return result;
}
summarize_tool_result(tool_name, arguments, &result)
}
/// Builds the MCP list section that replaces the `__MCP_LIST__` sentinel.
/// Resolves an `inject_memory` entry to `(absolute path to read, path to show)`.
///
/// `$WD` expands to the session's effective working directory (RunContext WD, or the
/// process cwd when unset). The shown path is **relative to that working directory
/// when the file lives under it, absolute otherwise** — so when the agent references
/// it back via `edit_file`/`write_file`, the loop's working-directory injection
/// (which rewrites relative paths against the WD) resolves to the very same file.
fn resolve_memory_path(&self, mem_path: &str) -> (std::path::PathBuf, String) {
let wd = self.working_directory.clone()
.unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
let expanded = mem_path.replace("$WD", &wd.display().to_string());
let abs = crate::tools::fs::resolve(&expanded)
.unwrap_or_else(|_| std::path::PathBuf::from(&expanded));
let display = match abs.strip_prefix(&wd) {
Ok(rel) => rel.to_string_lossy().into_owned(),
Err(_) => abs.to_string_lossy().into_owned(),
};
(abs, display)
}
fn render_mcp_list(&self, active_mcp_grants: &HashSet<String>) -> String {
let all_servers: std::collections::BTreeSet<String> = self.mcp.tools()
.into_iter()
.map(|t| t.server_name)
.collect();
if all_servers.is_empty() {
return String::new();
}
let descriptions = self.mcp.server_descriptions();
let hidden: Vec<&String> = all_servers.iter()
.filter(|n| !active_mcp_grants.contains(*n))
.collect();
let active: Vec<&String> = all_servers.iter()
.filter(|n| active_mcp_grants.contains(*n))
.collect();
let mut out = String::from("## MCP servers\n");
if !hidden.is_empty() {
out.push_str("\n**Available** — call `activate_tools([\"name\"])` to load tools:\n\n");
out.push_str("| Server | Description |\n|--------|-------------|\n");
for name in &hidden {
let desc = descriptions.get(*name)
.and_then(|d| d.as_deref())
.unwrap_or("");
out.push_str(&format!("| `{name}` | {desc} |\n"));
}
}
if !active.is_empty() {
out.push_str("\n**Active** — tools callable as `mcp__<name>__<tool>`:\n");
for name in &active {
out.push_str(&format!("- `{name}`\n"));
}
}
out
}
}
// ── Free helpers ──────────────────────────────────────────────────────────────
/// Creates an informative 1-line summary of a tool call result.
///
/// Produces human-readable descriptions like:
/// ```text
/// [execute_cmd] ran `cargo build` → exit 0, 47 lines output
/// [read_file] read src/main.rs (3,200 chars)
/// [write_file] wrote to agents/foo/AGENT.md
/// ```
fn summarize_tool_result(tool_name: &str, arguments: Option<&str>, result: &str) -> String {
let args: serde_json::Value = arguments
.and_then(|a| serde_json::from_str(a).ok())
.unwrap_or(serde_json::Value::Null);
let char_count = result.len();
let line_count = if result.trim().is_empty() { 0 } else { result.lines().count() };
fn arg_str<'a>(args: &'a serde_json::Value, key: &str) -> &'a str {
args[key].as_str().unwrap_or("?")
}
match tool_name {
tn::EXECUTE_CMD => {
let cmd = args["command"].as_str().unwrap_or("");
let cmd_display = super::preview_truncate(cmd, 77);
let exit_code = result
.lines()
.next()
.and_then(|l| l.strip_prefix("exit: "))
.unwrap_or("?");
format!("[execute_cmd] ran `{cmd_display}` → exit {exit_code}, {line_count} lines output")
}
"read_file" | "read_file_chunk" => {
let path = arg_str(&args, "path");
format!("[{tool_name}] read {path} ({char_count} chars)")
}
"write_file" => {
let path = arg_str(&args, "path");
format!("[write_file] wrote to {path}")
}
"edit_file" | "patch_file" => {
let path = arg_str(&args, "path");
format!("[{tool_name}] edited {path}")
}
"list_dir" | "glob" => {
let path = args["path"].as_str()
.or_else(|| args["pattern"].as_str())
.unwrap_or("?");
format!("[{tool_name}] {path} ({char_count} chars)")
}
"list_items" => {
let kind = arg_str(&args, "type");
format!("[list_items] {kind} ({char_count} chars)")
}
"toggle_item" => {
let kind = arg_str(&args, "kind");
let id = arg_str(&args, "id");
let enabled = args["enabled"].as_bool().unwrap_or(false);
format!("[toggle_item] {kind} '{id}' → {}", if enabled { "enabled" } else { "disabled" })
}
tn::READ_NOTIFICATION => {
let count = serde_json::from_str::<Vec<serde_json::Value>>(result)
.map(|v| v.len())
.unwrap_or(0);
format!("[read_notification] {count} notification(s)")
}
tn::EXECUTE_TASK | tn::EXECUTE_SUBTASK => {
let agent = arg_str(&args, "agent_id");
format!("[{tool_name}] → {agent} ({char_count} chars result)")
}
tn::ACTIVATE_TOOLS => {
let groups = args["groups"]
.as_array()
.map(|a| a.iter().filter_map(|v| v.as_str()).collect::<Vec<_>>().join(", "))
.unwrap_or_else(|| "?".to_string());
format!("[activate_tools] loaded: {groups}")
}
_ if tool_name.starts_with("mcp__") => {
format!("[{tool_name}] ({char_count} chars result)")
}
_ => {
let first_arg = args.as_object()
.and_then(|m| m.iter().next())
.map(|(k, v)| {
let sv = super::preview_truncate(v.as_str().unwrap_or_default(), 40);
format!(" {k}={sv}")
})
.unwrap_or_default();
format!("[{tool_name}]{first_arg} ({char_count} chars result)")
}
}
}
@@ -0,0 +1,44 @@
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use serde_json::Value;
use super::ChatSessionHandler;
use super::message_builder::MessageBuilder;
impl ChatSessionHandler {
/// Thin wrapper: constructs a `MessageBuilder` from this handler's fields
/// and delegates to `MessageBuilder::build`.
///
/// See `MessageBuilder::build` for the full documentation and message ordering.
pub(super) async fn build_openai_messages(
&self,
pool: &sqlx::SqlitePool,
stack_id: i64,
agent_id: &str,
extra_system_static: Option<&str>,
extra_system_dynamic: Option<&str>,
tail_reminder: Option<&str>,
active_mcp_grants: &HashSet<String>,
system_substitutions: &HashMap<String, String>,
cache_hints: bool,
) -> anyhow::Result<Vec<Value>> {
let effective_wd = self.run_context.read().await
.as_ref()
.map(|rc| rc.effective_working_dir());
let builder = MessageBuilder {
pool: Arc::clone(&self.db),
session_id: self.scratchpad_sid(),
mcp: Arc::clone(&self.mcp),
datetime_config: self.datetime_config.clone(),
max_history_messages: self.max_history_messages,
max_tool_result_chars: self.max_tool_result_chars,
compactor: self.compactor.clone(),
working_directory: effective_wd,
};
// `pool` is passed in from the caller (always `&self.db`) but we take
// ownership via Arc::clone above so the signature stays backward-compatible.
let _ = pool; // suppress unused-variable warning; MessageBuilder uses its own Arc
builder.build(stack_id, agent_id, extra_system_static, extra_system_dynamic, tail_reminder, active_mcp_grants, system_substitutions, cache_hints).await
}
}
@@ -0,0 +1,677 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use async_trait::async_trait;
use serde_json::{Value, json};
use sqlx::SqlitePool;
use tokio::sync::{Mutex, mpsc};
use tokio_util::sync::CancellationToken;
use tracing::{error, info, trace, warn};
use crate::approval::ApprovalManager;
use crate::run_context::RunContext;
use crate::tools::tool_names as tn;
use crate::chat_event_bus::{ChatEvent, ChatEventBus, ChatEventRole};
use crate::clarification::ClarificationManager;
use crate::compactor::ContextCompactor;
use crate::config::DatetimeConfig;
use crate::db::{chat_history, chat_sessions_stack};
use crate::events::ServerEvent;
use core_api::message_meta::MessageMetadata;
use crate::llm::LlmManager;
use crate::mcp::McpManager;
use crate::image_generate::ImageGeneratorManager;
use crate::memory::MemoryManager;
use crate::tool_discovery::ToolDiscovery;
use crate::tools::ToolRegistry;
mod approval;
mod agent_dispatch;
mod config;
mod dispatch;
mod emitter;
mod gate;
mod interface_tools;
mod llm_call;
mod llm_loop;
pub mod message_builder;
mod messages;
mod outcome;
mod resume;
use emitter::TurnEmitter;
pub use interface_tools::{InterfaceTool, ToolFuture};
pub const DEFAULT_MAX_TOOL_ROUNDS: usize = 20;
/// Default maximum number of synchronous sub-agents dispatched concurrently when
/// the LLM emits a homogeneous batch of sub-agent calls in a single response.
/// Bounds fan-out so a large batch does not trigger provider rate-limit storms.
pub const DEFAULT_MAX_PARALLEL_SUBAGENTS: usize = 4;
pub(super) const MAX_AGENT_DEPTH: i64 = 5;
/// A queued user message to be appended to history mid-turn (drained from the
/// source inbox at a round boundary).
pub struct PendingMsg {
pub content: String,
pub metadata: Option<MessageMetadata>,
}
/// Source of queued user input for the in-flight turn. Implemented by `ChatHub`
/// over a source's inbox; it lets `run_agent_turn` pull newly-queued user
/// messages at each round boundary and inject them live into the running turn.
///
/// Passed as `Some` only for the root interactive turn. Sub-agents, resume, and
/// non-interactive runners (cron, TIC) pass `None` — they never inject.
#[async_trait]
pub trait PendingUserInput: Send + Sync {
/// Drains the leading run of queued non-synthetic user messages, one entry
/// each. Returns empty when there is nothing to inject.
async fn drain_user(&self) -> Vec<PendingMsg>;
}
/// Control-flow signals returned as `anyhow::Error` by internal dispatch methods.
/// Using a typed enum instead of two separate sentinel structs allows a single
/// `downcast_ref` in `llm_loop` instead of two separate type checks.
#[derive(Debug)]
pub(super) enum AgentFlowSignal {
/// The WS disconnected while `dispatch_ask_user_clarification` was blocking.
/// The tool stays `'pending'` in DB so `resume_pending_tools` can re-ask on reconnect.
QuestionChannelClosed,
}
impl std::fmt::Display for AgentFlowSignal {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::QuestionChannelClosed => write!(f, "question channel closed (WS disconnected)"),
}
}
}
impl std::error::Error for AgentFlowSignal {}
pub(super) enum TurnOutcome {
Final {
content: String,
message_id: i64,
input_tokens: Option<u32>,
output_tokens: Option<u32>,
truncated: bool,
/// All tool calls executed during this turn, across all rounds.
tool_calls: Vec<crate::chat_event_bus::ToolCallEvent>,
},
Cancelled,
Exhausted,
}
/// Truncate `s` to at most `max_chars` characters, appending `…` when it was
/// longer. Char-boundary safe: a raw `&s[..n]` byte slice panics when byte `n`
/// lands inside a multi-byte UTF-8 character (e.g. an em-dash or emoji straddling
/// the cut point), which is exactly how a well-formed sub-agent result once
/// unwound a whole turn. Used for every event/log preview.
pub(super) fn preview_truncate(s: &str, max_chars: usize) -> String {
match s.char_indices().nth(max_chars) {
Some((byte_idx, _)) => format!("{}", &s[..byte_idx]),
None => s.to_string(),
}
}
pub(super) fn update_scratchpad_tool_def() -> Value {
json!({
"type": "function",
"function": {
"name": tn::UPDATE_SCRATCHPAD,
"description": "Write or update a key-value note in the session scratchpad. \
Notes are shared by all agents in this chat session and automatically \
injected into every agent's context. Not persisted across sessions. \
Use it for temporary discoveries: architecture notes, path lookups, \
decisions that other agents in this session need to know about.",
"parameters": {
"type": "object",
"properties": {
"key": { "type": "string", "description": "Short identifier for this note (e.g. 'db_url', 'main_struct')." },
"value": { "type": "string", "description": "Content of the note." }
},
"required": ["key", "value"]
}
}
})
}
/// Tool definition for `write_todos` — a private, per-turn task list the agent
/// uses to plan and track its own progress.
///
/// Unlike `update_scratchpad` (a shared blackboard injected into every agent in
/// the session), `write_todos` is **stateless**: the list lives only in this
/// agent's own tool-result history. Because conversation history is per-stack,
/// it is never visible to sub-agents or to the caller — no DB storage needed.
/// The agent re-sends the whole list (TodoWrite-style) on every update.
pub(super) fn write_todos_tool_def() -> Value {
json!({
"type": "function",
"function": {
"name": tn::WRITE_TODOS,
"description": "Record and update your task list for the current turn, to plan multi-step \
work and track progress. Re-send the ENTIRE list on every call (including \
already-completed items with their new status) — this replaces the previous \
list. Keep exactly one item `in_progress` at a time. This list is PRIVATE \
to you: it is not shared with sub-agents you dispatch, nor returned to your \
caller (use `update_scratchpad` instead for notes other agents must see).",
"parameters": {
"type": "object",
"properties": {
"todos": {
"type": "array",
"description": "The full, ordered task list. Re-send it entirely on every update.",
"items": {
"type": "object",
"properties": {
"content": { "type": "string", "description": "Short description of the task." },
"status": { "type": "string", "enum": ["pending", "in_progress", "completed"], "description": "Current status of this task." }
},
"required": ["content", "status"]
}
}
},
"required": ["todos"]
}
}
})
}
/// Tool definition that lets a sub-agent (depth > 0) dispatch a further
/// synchronous sub-agent. The call is intercepted in `run_agent_turn` and routed
/// to `dispatch_sub_agent` (the InterfaceTool handler is never reached), so only
/// the definition is needed here. `agent_id` is required because
/// `dispatch_sub_agent` rejects calls without it.
fn execute_subtask_tool_def() -> Value {
json!({
"type": "function",
"function": {
"name": tn::EXECUTE_SUBTASK,
"description": "Delegate work to another agent and get its result. Runs the \
named agent synchronously with the given prompt and blocks until \
it finishes, returning its final answer as the tool result. Use \
`list_agents` first to see which agents are available.",
"parameters": {
"type": "object",
"properties": {
"agent_id": { "type": "string", "description": "Id of the agent to run (see `list_agents`)." },
"title": { "type": "string", "description": "Short name for this sub-task." },
"description": { "type": "string", "description": "What this sub-task does." },
"prompt": { "type": "string", "description": "Prompt sent to the agent." }
},
"required": ["agent_id", "prompt"]
}
}
})
}
fn ask_user_clarification_tool_def() -> Value {
json!({
"type": "function",
"function": {
"name": tn::ASK_USER_CLARIFICATION,
"description": "Pause execution and ask the user a clarification question. \
Use when requirements are ambiguous, a dependency is missing, \
or a decision requires user input before continuing. \
The user's answer is returned as the tool result.",
"parameters": {
"type": "object",
"properties": {
"title": { "type": "string", "description": "Short label shown in the inbox card (e.g. 'Missing API key')." },
"question": { "type": "string", "description": "Full question text." },
"suggested_answers": {
"type": "array",
"items": { "type": "string" },
"description": "Optional list of suggested answers shown as chips. The user can pick one or type freely."
}
},
"required": ["title", "question"]
}
}
})
}
pub enum ApprovalDecision {
Approved,
Rejected { note: String },
}
impl ApprovalDecision {
/// Canonical tool-result text shown to the LLM for a human rejection,
/// given the raw user-supplied note (which may be empty). This is the
/// single source of truth: every reject path passes the raw note and lets
/// this build the message, so the wording stays consistent and the note
/// carries the user's justification verbatim — no surface-specific prefixes.
pub fn rejection_message(note: &str) -> String {
let note = note.trim();
if note.is_empty() {
"User rejected this tool call.".to_string()
} else {
format!("User rejected this tool call. Reason: {note}")
}
}
}
pub struct ChatSessionHandler {
pub session_id: i64,
pub(super) db: Arc<SqlitePool>,
pub(super) llm_manager: Arc<LlmManager>,
pub(super) max_history_messages: usize,
pub(super) max_tool_rounds: usize,
/// Max synchronous sub-agents dispatched concurrently for a homogeneous batch
/// of sub-agent calls in a single LLM response (`1` = sequential).
pub(super) max_parallel_subagents: usize,
/// If `Some(n)`, tool results from previous turns that exceed `n` characters
/// are replaced with a placeholder when building the LLM context.
/// The database always retains the original content.
pub(super) max_tool_result_chars: Option<usize>,
pub(super) datetime_config: DatetimeConfig,
pub(super) agent_id: String,
/// Source of the session: "web", "telegram", "cron", etc.
pub(super) source: String,
/// True when a real user is actively participating (web, telegram).
pub(super) is_interactive: bool,
/// True for short-lived automated sessions (cron, tic).
pub(super) is_ephemeral: bool,
pub(super) tools: Arc<ToolRegistry>,
pub(super) mcp: Arc<McpManager>,
/// Records tools offered to the LLM each round so the Security-groups UI can
/// list/gate dynamically-injected tools (interface/plugin/provider tools).
pub(super) tool_discovery: Arc<ToolDiscovery>,
pub(super) approval: Arc<ApprovalManager>,
pub(super) clarification: Arc<ClarificationManager>,
pub(super) event_bus: Arc<ChatEventBus>,
/// Human-readable label injected by background runners (e.g. "CronJob: Daily Digest").
pub(super) context_label: std::sync::RwLock<Option<String>>,
pub(super) memory_manager: Arc<MemoryManager>,
pub(super) image_generator_manager: Arc<ImageGeneratorManager>,
/// Prevents concurrent handle_message calls on the same session.
pub(super) processing: Mutex<()>,
/// Cancellation scope for the in-flight turn. A fresh token is minted per
/// user message (`handle_message`) and per resume (`resume_turn`), then a
/// clone is threaded by value through the whole (possibly recursive) call
/// tree. `cancel()` cancels whatever token is currently stored, which the
/// running chain observes because it holds its own clone of that same token.
/// Replacing the field only affects the *next* turn — that is what makes a
/// stop sticky across sub-agent recursion (it is never reset mid-turn).
pub(super) current_cancel: std::sync::Mutex<CancellationToken>,
/// When true, any tool call that would require human approval is automatically
/// denied instead of blocking. Used by TicManager and other headless runners
/// that cannot process approval requests.
pub(super) auto_deny_approvals: AtomicBool,
/// Tool-call ids the user already approved via a resolve endpoint after a restart
/// (no live oneshot to unblock). The next resume's approval gate skips re-gating
/// these so a post-restart approve dispatches the tool without a second prompt.
pub(super) pre_approved: std::sync::Mutex<std::collections::HashSet<i64>>,
/// Context compactor, shared across all sessions. `None` when compaction
/// is disabled (no `compaction` section in config).
pub(super) compactor: Option<Arc<ContextCompactor>>,
/// Input token count from the most recently completed turn, stored
/// atomically so the next `handle_message` call can decide whether to
/// compact before processing the new message. Zero means unknown
/// (provider did not report usage on the first turn).
pub(super) last_input_tokens: AtomicU32,
/// Active RunContext for this session. `None` means the "default" group is used implicitly.
pub(super) run_context: tokio::sync::RwLock<Option<RunContext>>,
/// When set, scratchpad reads/writes use this session_id instead of `self.session_id`.
/// Used by async sub-tasks to share the parent's scratchpad.
pub(super) scratchpad_session_id: std::sync::OnceLock<i64>,
}
impl ChatSessionHandler {
pub fn new(
session_id: i64,
db: Arc<SqlitePool>,
llm_manager: Arc<LlmManager>,
max_history_messages: usize,
max_tool_rounds: usize,
max_parallel_subagents: usize,
max_tool_result_chars: Option<usize>,
datetime_config: DatetimeConfig,
agent_id: String,
source: String,
is_interactive: bool,
is_ephemeral: bool,
tools: Arc<ToolRegistry>,
mcp: Arc<McpManager>,
approval: Arc<ApprovalManager>,
clarification: Arc<ClarificationManager>,
event_bus: Arc<ChatEventBus>,
memory_manager: Arc<MemoryManager>,
image_generator_manager: Arc<ImageGeneratorManager>,
compactor: Option<Arc<ContextCompactor>>,
run_context: Option<RunContext>,
tool_discovery: Arc<ToolDiscovery>,
) -> Self {
Self {
session_id,
db,
llm_manager,
max_history_messages,
max_tool_rounds,
max_parallel_subagents,
max_tool_result_chars,
datetime_config,
agent_id,
source,
is_interactive,
is_ephemeral,
tools,
mcp,
tool_discovery,
approval,
clarification,
event_bus,
memory_manager,
image_generator_manager,
compactor,
context_label: std::sync::RwLock::new(None),
processing: Mutex::new(()),
current_cancel: std::sync::Mutex::new(CancellationToken::new()),
auto_deny_approvals: AtomicBool::new(false),
pre_approved: std::sync::Mutex::new(std::collections::HashSet::new()),
last_input_tokens: AtomicU32::new(0),
run_context: tokio::sync::RwLock::new(run_context),
scratchpad_session_id: std::sync::OnceLock::new(),
}
}
/// Sets the human-readable context label for this session (e.g. "CronJob: Daily Digest").
/// Called by background runners after the handler is created.
pub fn set_context_label(&self, label: impl Into<String>) {
if let Ok(mut g) = self.context_label.write() {
*g = Some(label.into());
}
}
/// Override the session used for scratchpad reads/writes.
/// Called by the cron runner for async tasks so they share the parent's scratchpad.
pub fn set_scratchpad_session_id(&self, id: i64) {
let _ = self.scratchpad_session_id.set(id);
}
/// Returns the session_id to use for scratchpad operations.
pub(super) fn scratchpad_sid(&self) -> i64 {
*self.scratchpad_session_id.get().unwrap_or(&self.session_id)
}
/// Updates the active RunContext for this session at runtime.
pub async fn set_run_context(&self, ctx: Option<RunContext>) {
*self.run_context.write().await = ctx;
}
/// Returns the serialised JSON blob of the active RunContext (for storing on child tasks).
pub async fn run_context_json(&self) -> Option<String> {
self.run_context.read().await.as_ref().map(|rc| rc.to_db())
}
/// Returns the active tool_permission_groups id for approval checks.
pub(super) async fn tool_group_id(&self) -> Option<String> {
self.run_context.read().await.as_ref().and_then(|rc| rc.tool_group_id().map(str::to_owned))
}
/// Cancels the in-flight turn. The running call tree holds its own clone of
/// the same token, so it stops at the next round boundary, on the in-flight
/// LLM call, and on cancellable tools (e.g. `execute_cmd`). Sticky across
/// sub-agent recursion: the token is never reset mid-turn.
pub fn cancel(&self) {
self.current_cancel.lock().unwrap().cancel();
}
/// True if a turn is currently in flight (the `processing` mutex is held for
/// the whole duration of `handle_message` / `resume_turn`). Used to tell a
/// freshly (re)connected client to show the STOP button.
pub fn is_processing(&self) -> bool {
self.processing.try_lock().is_err()
}
/// When set, any tool call that would require human approval is automatically
/// denied instead of blocking indefinitely.
pub fn set_auto_deny_approvals(&self) {
self.auto_deny_approvals.store(true, Ordering::Relaxed);
}
/// Records that the user already approved this tool_call via a resolve endpoint
/// after a restart (no live oneshot to unblock). The next resume's approval gate
/// consumes this and skips re-gating, so the tool dispatches without re-prompting.
pub fn mark_pre_approved(&self, tool_call_id: i64) {
self.pre_approved.lock().unwrap().insert(tool_call_id);
}
/// Cancels all pending approvals for this session in the ApprovalManager.
/// Called when the WS connection is lost mid-approval so the waiting future unblocks.
pub async fn cancel_pending_approvals(&self) {
self.approval.cancel_for_session(self.session_id).await;
}
/// Resolves a pending `ask_user_clarification` call with the user's answer.
pub async fn resolve_question(&self, request_id: i64, answer: String) {
if !self.clarification.resolve(request_id, answer).await {
warn!(session_id = self.session_id, request_id, "resolve_question: request_id not found in ClarificationManager");
}
}
/// Cancels all pending clarification requests for this session (WS disconnected).
/// The blocked `rx.await` in dispatch_ask_user_clarification returns Err → TurnOutcome::Cancelled,
/// leaving the tool as 'pending' so resume_pending_tools re-dispatches on reconnect.
pub async fn cancel_pending_questions(&self) {
self.clarification.cancel_for_session(self.session_id).await;
}
/// Force compaction of the current stack's conversation history.
/// Bypasses the token threshold check; still respects the ephemeral guard.
/// Returns `true` if a new summary was written, `false` if skipped.
pub async fn force_compact(&self) -> anyhow::Result<bool> {
let pool = &self.db;
let stack = match chat_sessions_stack::active_for_session(pool, self.session_id).await? {
Some(s) => s,
None => return Ok(false),
};
match self.compactor {
Some(ref compactor) => {
compactor.force_compact(pool, self.session_id, stack.id, self.is_ephemeral).await
}
None => Ok(false),
}
}
/// Processes a user message end-to-end:
/// saves it, runs the tool-calling loop, saves the final response,
/// sends a Done event. Only one call can run at a time per session.
pub async fn handle_message(
&self,
content: &str,
client_name: Option<String>,
extra_system_context: Option<String>,
// Per-turn dynamic system suffix injected AFTER conversation history.
// Merged with the Honcho memory context (which also lives at position 5).
// Use for per-turn framing that must not pollute the cacheable static prefix
// (e.g. notification behavioural instructions from ChatHub).
extra_system_dynamic_override: Option<String>,
tail_reminder: Option<String>,
interface_tools: Vec<InterfaceTool>,
system_substitutions: HashMap<String, String>,
tx: mpsc::Sender<ServerEvent>,
// True for system-generated messages injected as user turns
// (TicManager ticks, notification briefings from ChatHub).
is_synthetic: bool,
// Structured metadata persisted on the user turn (e.g. file attachments).
// The MessageBuilder derives the LLM-facing block; the UI renders chips.
metadata: Option<MessageMetadata>,
// Queued user input for this source. When `Some`, `run_agent_turn` drains
// it at each round boundary and injects newly-arrived user messages into
// the running turn. `None` for sub-agents / resume / non-interactive runners.
pending_input: Option<Arc<dyn PendingUserInput>>,
) -> anyhow::Result<()> {
let _guard = self.processing.lock().await;
// Fresh cancellation scope for this user message. Stored so `cancel()`
// can reach it, and cloned-by-value into the call tree so a /stop during
// the turn is sticky across sub-agent recursion (never reset mid-turn).
let token = CancellationToken::new();
*self.current_cancel.lock().unwrap() = token.clone();
let pool = &self.db;
let em = TurnEmitter::new(&tx);
// Retrieve memory context (Honcho or other backend) for this turn.
// Kept SEPARATE from extra_system_context (the static part) so it can be
// injected as a dynamic tail system message after the conversation history
// rather than embedded in the cacheable static prefix. This allows
// providers with prefix caching (e.g. Alibaba/DeepSeek via OpenRouter)
// to cache the stable system prompt across turns even though Honcho
// memories change on every call.
let honcho_dynamic = match self.memory_manager.query_context(self.session_id, content).await {
Some(mem_ctx) => {
trace!(
session_id = self.session_id,
chars = mem_ctx.len(),
"handle_message: memory context retrieved (will be injected as dynamic tail)"
);
Some(mem_ctx)
}
None => {
trace!(
session_id = self.session_id,
"handle_message: no memory context returned (cold start, unavailable, or nothing to say)"
);
None
}
};
// Merge Honcho memories with any per-turn override from the caller.
// The override goes last so it sits closest to the generation point (recency bias).
// extra_system_context (passed by the caller) is the STATIC part:
// interface-specific formatting rules (e.g. Telegram HTML format),
// never changes turn-to-turn, safe to include in the cached prefix.
let extra_system_dynamic = match (honcho_dynamic, extra_system_dynamic_override) {
(Some(honcho), Some(override_)) => Some(format!("{honcho}\n\n{override_}")),
(Some(honcho), None) => Some(honcho),
(None, Some(override_)) => Some(override_),
(None, None) => None,
};
let mut config = self.build_agent_config(
client_name, extra_system_context, extra_system_dynamic, interface_tools, system_substitutions,
).await?;
config.tail_reminder = tail_reminder;
let stack = match chat_sessions_stack::active_for_session(pool, self.session_id).await? {
Some(s) => s,
None => {
chat_sessions_stack::create(pool, self.session_id, "main", None, 0, None).await?
}
};
info!(session_id = self.session_id, stack_id = stack.id, client = %config.client_name, "handle_message start");
// ── Context compaction (Opzione C: at the start of the next turn) ────
// Check whether the previous turn's input token count exceeded the
// threshold. If so, summarise the old history before processing the
// new message. This keeps latency transparent to the user — the wait
// happens here, before the LLM loop, and is not a separate turn.
if let Some(ref compactor) = self.compactor {
let last_tokens = self.last_input_tokens.load(Ordering::Relaxed);
match compactor.try_compact(pool, self.session_id, stack.id, last_tokens, self.is_ephemeral).await {
Ok(true) => info!(session_id = self.session_id, stack_id = stack.id, "handle_message: context compacted"),
Ok(false) => {}
Err(e) => warn!(session_id = self.session_id, error = %e, "handle_message: compaction failed (non-fatal), continuing"),
}
}
// ─────────────────────────────────────────────────────────────────────
// If the previous turn was cancelled before the LLM responded, the history ends on a
// User message with no following assistant. This breaks the user→assistant alternation
// required by strict APIs (e.g. OpenRouter). Mark the orphaned message as failed so
// for_stack() excludes it from the context we send to the LLM.
let prior = chat_history::for_stack(pool, stack.id).await?;
if let Some(last) = prior.last() {
if matches!(last.role, chat_history::Role::User | chat_history::Role::Agent) {
warn!(session_id = self.session_id, message_id = last.id, "orphaned user message (cancelled turn) — marking failed");
chat_history::mark_failed(pool, last.id).await?;
}
}
let user_content = content.to_string(); // save before TurnOutcome::Final shadows `content`
let user_message_id = chat_history::append_with_metadata(pool, stack.id, &chat_history::Role::User, content, is_synthetic, None, metadata.as_ref()).await?;
// Telnet-style echo: the bubble appears only once the message is persisted.
// Synthetic turns (TIC/notification) never produce a user bubble.
if !is_synthetic {
let attachments = metadata.as_ref().map(|m| m.attachments.clone()).unwrap_or_default();
// A custom slash command persists its expanded template (for LLM replay)
// but the bubble must show the typed command — emit `display` when present.
let echo = metadata.as_ref()
.and_then(|m| m.command.as_ref())
.map(|c| c.display.clone())
.unwrap_or_else(|| user_content.clone());
em.user_message(user_message_id, echo, attachments).await;
}
// Resume any tool calls left pending from a previous interrupted session.
// They are re-gated (rules may have changed) and executed before the LLM runs.
self.resume_pending_tools(stack.id, &config, &token, &tx).await?;
let outcome = self.run_agent_turn(stack.id, &config, &token, &tx, pending_input.as_ref()).await?;
match outcome {
TurnOutcome::Final { content, message_id, input_tokens, output_tokens, truncated, tool_calls } => {
// Persist token count so the *next* handle_message call knows
// whether to compact before running the LLM loop.
if let Some(t) = input_tokens {
self.last_input_tokens.store(t, Ordering::Relaxed);
}
info!(session_id = self.session_id, stack_id = stack.id, ?input_tokens, ?output_tokens, "handle_message done");
if truncated {
warn!(session_id = self.session_id, ?output_tokens, "response truncated (max_tokens)");
em.truncated(output_tokens).await;
}
em.done(message_id, stack.id, content.clone(), input_tokens, output_tokens).await;
// Publish both messages to the event bus now that both are in the DB.
let now = chrono::Utc::now();
self.event_bus.user_message(ChatEvent {
session_id: self.session_id,
stack_id: stack.id,
message_id: user_message_id,
role: ChatEventRole::User,
content: user_content,
is_synthetic,
is_interactive: self.is_interactive,
is_ephemeral: self.is_ephemeral,
tool_calls: vec![],
created_at: now,
});
self.event_bus.assistant_response(ChatEvent {
session_id: self.session_id,
stack_id: stack.id,
message_id,
role: ChatEventRole::Assistant,
content,
is_synthetic: false,
is_interactive: self.is_interactive,
is_ephemeral: self.is_ephemeral,
tool_calls,
created_at: now,
});
Ok(())
}
TurnOutcome::Cancelled => {
info!(session_id = self.session_id, "handle_message cancelled by user");
em.error("Cancelled by user.".to_string()).await;
Err(anyhow::anyhow!("Turn cancelled by user"))
}
TurnOutcome::Exhausted => {
error!(session_id = self.session_id, max_rounds = self.max_tool_rounds, "tool-call loop exhausted without final answer");
em.error(format!("Exceeded {} tool-call rounds without a final answer.", self.max_tool_rounds)).await;
Err(anyhow::anyhow!("tool-call loop exhausted after {} rounds without a final answer", self.max_tool_rounds))
}
}
}
}
@@ -0,0 +1,92 @@
//! Shared recording of a single tool-call outcome.
//!
//! The persist-then-emit tail of a tool call (`ExecutionOutcome` → DB row +
//! `ToolDone`/`ToolError`/`ToolCancelled` event) was copy-pasted into both the live
//! loop (`run_agent_turn`) and `resume_pending_tools`. `record_tool_outcome` is the
//! single implementation both call.
use serde_json::Value;
use tracing::{debug, info, warn};
use crate::chat_event_bus::ToolCallEvent;
use crate::db::chat_llm_tools;
use crate::tools::{is_file_write_tool, ExecutionOutcome};
use super::ChatSessionHandler;
use super::emitter::TurnEmitter;
/// Whether the enclosing loop should keep going after an outcome is recorded.
pub(super) enum RecordFlow {
/// Continue with the next tool call / round.
Continue,
/// The tool was cancelled by the user — the caller must end the turn.
Abort,
}
impl ChatSessionHandler {
/// Persists one tool-call outcome and emits the matching lifecycle event.
/// Returns [`RecordFlow::Abort`] for a user cancellation (the caller ends the
/// turn), [`RecordFlow::Continue`] otherwise.
///
/// When `accumulate` is `Some` (the live turn), the call is also appended to the
/// turn's `ToolCallEvent` list for the chat-event bus, and a `FileChanged` event
/// is emitted for a successful file-write tool. `resume_pending_tools` passes
/// `None`: it neither accumulates nor re-emits `FileChanged`.
pub(super) async fn record_tool_outcome(
&self,
tool_call_id: i64,
tool_name: &str,
args: &Value,
outcome: ExecutionOutcome,
em: &TurnEmitter<'_>,
accumulate: Option<&mut Vec<ToolCallEvent>>,
) -> anyhow::Result<RecordFlow> {
let pool = &self.db;
match outcome {
ExecutionOutcome::Completed(result) => {
let wire = result.to_wire();
let kind = result.kind();
debug!(session_id = self.session_id, tool = %tool_name, tool_call_id, result_len = wire.len(), "tool done");
chat_llm_tools::complete(pool, tool_call_id, &wire, kind).await?;
if let Some(acc) = accumulate {
if is_file_write_tool(tool_name)
&& let Some(p) = args["path"].as_str()
{
em.file_changed(crate::approval::normalize_path(p)).await;
}
acc.push(ToolCallEvent {
name: tool_name.to_string(),
arguments: Some(serde_json::to_string(args).unwrap_or_default()),
result: Some(wire.clone()),
status: "done".to_string(),
});
}
em.tool_done(tool_call_id, wire, kind.to_string()).await;
Ok(RecordFlow::Continue)
}
ExecutionOutcome::Failed(msg) => {
warn!(session_id = self.session_id, tool = %tool_name, tool_call_id, error = %msg, "tool failed");
chat_llm_tools::fail(pool, tool_call_id, &msg).await?;
if let Some(acc) = accumulate {
acc.push(ToolCallEvent {
name: tool_name.to_string(),
arguments: Some(serde_json::to_string(args).unwrap_or_default()),
result: Some(msg.clone()),
status: "failed".to_string(),
});
}
em.tool_error(tool_call_id, msg).await;
Ok(RecordFlow::Continue)
}
ExecutionOutcome::Cancelled => {
// A /stop hit this tool mid-flight. Record it as cancelled (not
// failed); the sticky token cancels the rest of the loop by
// construction, so the caller just ends the turn.
info!(session_id = self.session_id, tool = %tool_name, tool_call_id, "tool cancelled by user");
chat_llm_tools::cancel(pool, tool_call_id, "Cancelled by user.").await?;
em.tool_cancelled(tool_call_id).await;
Ok(RecordFlow::Abort)
}
}
}
}
@@ -0,0 +1,401 @@
use serde_json::Value;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};
use crate::db::{chat_history, chat_llm_tools, chat_sessions_stack};
use crate::events::ServerEvent;
use crate::tools::{ToolDescriptionLength, ToolResult, tool_names as tn};
use super::{ChatSessionHandler, TurnOutcome};
use super::emitter::TurnEmitter;
use super::gate::GateOutcome;
use super::outcome::RecordFlow;
use super::interface_tools::{AgentRunConfig, InterfaceTool};
impl ChatSessionHandler {
/// Dispatches a single tool call by name+args without going through the LLM loop.
/// Used by the REST `resolve` endpoint and by `resume_pending_tools`.
/// Does NOT update the DB — caller is responsible for `complete` / `fail`.
pub async fn execute_tool(&self, name: &str, args: Value) -> anyhow::Result<ToolResult> {
if let Some((srv, mcp_tool)) = crate::mcp::parse_mcp_tool_name(name) {
return self.mcp.call(srv, mcp_tool, args).await;
}
self.tools.dispatch(name, args).await.map(ToolResult::Text)
}
/// Resumes the LLM loop for the current session WITHOUT appending a new user message.
/// Intended for use after pending tool calls have been resolved externally
/// (e.g. via the REST approve endpoint) so the LLM can produce a final response
/// or make further tool calls using the now-complete history.
pub async fn resume_turn(
&self,
client_name: Option<String>,
extra_system_context: Option<String>,
interface_tools: Vec<InterfaceTool>,
tx: mpsc::Sender<ServerEvent>,
) -> anyhow::Result<()> {
let _guard = self.processing.lock().await;
// A resume is a fresh unit of work (async result injection, app-restart
// recovery, WS resume): mint a new token so it does not inherit a stale
// cancellation, while a /stop *during* the resume still cancels this token.
let token = CancellationToken::new();
*self.current_cancel.lock().unwrap() = token.clone();
let pool = &self.db;
let em = TurnEmitter::new(&tx);
let mut config = self.build_agent_config(
client_name, extra_system_context, None, interface_tools, std::collections::HashMap::new(),
).await?;
config.tail_reminder = None;
// Prune any interrupted parallel sub-agent batch before the linear cascade,
// which assumes a single active frame per depth (see method doc).
self.reap_interrupted_parallel_batches().await?;
let stack = match chat_sessions_stack::active_for_session(pool, self.session_id).await? {
Some(s) => s,
None => {
warn!(session_id = self.session_id, "resume_turn: no active stack, nothing to resume");
return Ok(());
}
};
info!(session_id = self.session_id, stack_id = stack.id, depth = stack.depth, "resume_turn start");
// Resume pending/interrupted tools before running the LLM loop.
let had_pending = self.resume_pending_tools(stack.id, &config, &token, &tx).await?;
// Seed the cascade. Normally we (re)run the deepest active frame's LLM loop
// (live injection only applies to a fresh interactive turn from handle_message).
// Two special cases when nothing was pending AND the frame's last message is a
// pure-text assistant reply (its own turn is already complete):
// • root frame (no parent) → nothing to do, skip the LLM.
// • child frame (has parent) → its result was produced but never propagated
// (e.g. the turn task died right after the child finished). Seed the cascade
// from the existing final message — without re-running the LLM — so the
// parent's tool call is completed and the parent continues. Skipping here
// (as the old guard did unconditionally) left the parent wedged forever.
let (mut current_outcome, mut current_stack) = 'seed: {
if !had_pending {
if let Some(msg) = chat_history::last_message_for_stack(pool, stack.id).await? {
if matches!(msg.role, chat_history::Role::Assistant)
&& chat_llm_tools::for_message(pool, msg.id).await?.is_empty()
{
if stack.parent_tool_call_id.is_none() {
info!(session_id = self.session_id, stack_id = stack.id, "resume_turn: last message is pure-text assistant, turn already complete — skipping LLM");
return Ok(());
}
info!(session_id = self.session_id, stack_id = stack.id, "resume_turn: deepest frame is a completed child — cascading its existing result to the parent");
let outcome = TurnOutcome::Final {
content: msg.content,
message_id: msg.id,
input_tokens: None,
output_tokens: None,
truncated: false,
tool_calls: Vec::new(),
};
break 'seed (outcome, stack);
}
}
}
(self.run_agent_turn(stack.id, &config, &token, &tx, None).await?, stack)
};
// Cascade completion upward through parent stacks (handles app-restart recovery
// when a sub-agent was running — child completes, then parent continues).
loop {
let Some(parent_tool_call_id) = current_stack.parent_tool_call_id else { break };
// Determine the result string to propagate to the parent's call_agent tool.
let (result_str, is_error) = match &current_outcome {
TurnOutcome::Final { content, .. } => (content.clone(), false),
TurnOutcome::Cancelled => (format!("Sub-agent `{}` was cancelled.", current_stack.agent_id), true),
TurnOutcome::Exhausted => (format!("Sub-agent `{}` exhausted tool-call rounds.", current_stack.agent_id), true),
};
let result_preview = super::preview_truncate(&result_str, 500);
// Complete or fail the parent's call_agent tool call.
if is_error {
chat_llm_tools::fail(pool, parent_tool_call_id, &result_str).await?;
} else {
chat_llm_tools::complete(pool, parent_tool_call_id, &result_str, "string").await?;
}
// Terminate the child stack so active_for_session() returns the parent next.
let _ = chat_sessions_stack::terminate(pool, current_stack.id).await;
// Emit events to the frontend.
if is_error {
em.tool_error(parent_tool_call_id, result_str).await;
} else {
em.tool_done(parent_tool_call_id, result_str, "string".to_string()).await;
}
// Now the parent is the deepest active stack.
let parent_stack = match chat_sessions_stack::active_for_session(pool, self.session_id).await? {
Some(s) => s,
None => {
warn!(session_id = self.session_id, "resume_turn cascade: no active stack after child terminated");
break;
}
};
em.agent_done(
current_stack.id,
current_stack.agent_id.clone(),
parent_stack.agent_id.clone(),
result_preview,
).await;
info!(
session_id = self.session_id,
child_stack = current_stack.id,
parent_stack = parent_stack.id,
depth = parent_stack.depth,
"resume_turn: cascading to parent stack"
);
self.resume_pending_tools(parent_stack.id, &config, &token, &tx).await?;
current_outcome = self.run_agent_turn(parent_stack.id, &config, &token, &tx, None).await?;
current_stack = parent_stack;
}
// current_stack is now the root (depth=0); emit the final event.
match current_outcome {
TurnOutcome::Final { content, message_id, input_tokens, output_tokens, truncated, .. } => {
info!(session_id = self.session_id, "resume_turn done");
if truncated {
warn!(session_id = self.session_id, "response truncated");
em.truncated(output_tokens).await;
}
em.done(message_id, current_stack.id, content, input_tokens, output_tokens).await;
}
TurnOutcome::Cancelled => {
info!(session_id = self.session_id, "resume_turn cancelled");
em.error("Cancelled by user.".to_string()).await;
}
TurnOutcome::Exhausted => {
error!(session_id = self.session_id, "resume_turn exhausted tool rounds");
em.error("Exceeded tool-call rounds without a final answer.".to_string()).await;
}
}
Ok(())
}
/// Restart recovery for an interrupted **parallel** sub-agent batch.
///
/// A purely linear stack has at most one active frame per depth. Two or more
/// active frames at the same depth can only mean a concurrent sub-agent batch
/// (`handle_sub_agent_batch`) was in flight when the process died. This app is
/// single-user and deliberately tolerates losing mid-turn work on restart, so
/// rather than a complex multi-sibling re-drive we simply prune the batch:
/// terminate every active frame from the shallowest multi-frame depth downward
/// and fail the sub-agent tool call that spawned each. The parent frame is then
/// left with a clean, fully-resolved set of tool calls and the normal linear
/// cascade resumes it. A single interrupted sub-agent (one frame at its depth)
/// is untouched and still recovers via the existing cascade.
async fn reap_interrupted_parallel_batches(&self) -> anyhow::Result<()> {
let pool = &self.db;
let active = chat_sessions_stack::active_all_for_session(pool, self.session_id).await?;
let Some(d_min) = shallowest_parallel_depth(&active) else {
return Ok(()); // linear stack — nothing to reap
};
warn!(
session_id = self.session_id, depth = d_min,
"restart recovery: pruning interrupted parallel sub-agent batch"
);
for frame in active.iter().filter(|f| f.depth >= d_min) {
if let Some(parent_tool_call_id) = frame.parent_tool_call_id {
let _ = chat_llm_tools::fail(
pool, parent_tool_call_id, "Sub-agent interrupted by restart (parallel batch).",
).await;
}
let _ = chat_sessions_stack::terminate(pool, frame.id).await;
}
Ok(())
}
/// Called at the start of `handle_message` (and by the REST endpoint after a manual
/// resolve). Finds any `pending` tool calls left from a previous interrupted session,
/// re-runs them through the approval gate, executes approved ones, and fails rejected
/// or denied ones — so `run_agent_turn` sees complete history and can continue cleanly.
pub async fn resume_pending_tools(
&self,
stack_id: i64,
config: &AgentRunConfig,
token: &CancellationToken,
tx: &mpsc::Sender<ServerEvent>,
) -> anyhow::Result<bool> {
let pool = &self.db;
let em = TurnEmitter::new(tx);
let pending = chat_llm_tools::pending_for_stack(pool, stack_id).await?;
if pending.is_empty() {
return Ok(false);
}
info!(
session_id = self.session_id, stack_id,
count = pending.len(), "resuming pending tool calls"
);
for tc in pending {
let args: Value = tc.arguments.as_deref()
.and_then(|s| serde_json::from_str(s).ok())
.unwrap_or(Value::Object(Default::default()));
// A pending `execute_task` (mode=sync) or `execute_subtask` means a
// sub-agent stack was active. The cascade in resume_turn() handles it
// by running the child stack to completion and propagating the result
// up — skip it here.
if tc.name == tn::EXECUTE_TASK || tc.name == tn::EXECUTE_SUBTASK {
info!(session_id = self.session_id, tool_call_id = tc.id, "resume: skipping sub-agent dispatch (handled by stack cascade)");
continue;
}
// `ask_user_clarification` is a synthetic tool (not in the registry).
// Re-dispatch it directly so the question is re-asked to the user.
if tc.name == tn::ASK_USER_CLARIFICATION {
info!(session_id = self.session_id, tool_call_id = tc.id, "resume: re-asking clarification question");
em.tool_start(
tc.id,
tc.message_id,
tc.name.clone(),
args.clone(),
self.tools.describe_call(&tc.name, &args, ToolDescriptionLength::Short),
self.tools.describe_call(&tc.name, &args, ToolDescriptionLength::Full),
self.tools.target_path(&tc.name, &args),
).await;
let result = self.dispatch_ask_user_clarification(tc.id, &args, tx).await;
match result {
Ok(answer) => {
chat_llm_tools::complete(pool, tc.id, &answer, "string").await?;
em.tool_done(tc.id, answer, "string".to_string()).await;
}
Err(e) if matches!(e.downcast_ref::<super::AgentFlowSignal>(), Some(super::AgentFlowSignal::QuestionChannelClosed)) => {
// WS disconnected again mid-resume. Tool stays 'pending' — next resume re-asks.
warn!(session_id = self.session_id, tool_call_id = tc.id, "clarification channel closed during resume — aborting");
return Ok(true);
}
Err(e) => {
let msg = e.to_string();
chat_llm_tools::fail(pool, tc.id, &msg).await?;
em.tool_error(tc.id, msg).await;
}
}
continue;
}
// Announce the tool is being re-tried.
em.tool_start(
tc.id,
tc.message_id,
tc.name.clone(),
args.clone(),
self.tools.describe_call(&tc.name, &args, ToolDescriptionLength::Short),
self.tools.describe_call(&tc.name, &args, ToolDescriptionLength::Full),
self.tools.target_path(&tc.name, &args),
).await;
// Re-run through the same approval gate as a live turn (current rules,
// RunContext fast-path, auto-deny). Deny/reject paths mark the DB row and
// emit the event internally; a closed channel leaves the tool pending.
match self.run_approval_gate(tc.id, &tc.name, &args, &config.agent_id, &em).await? {
GateOutcome::Proceed => {}
GateOutcome::Rejected => continue,
GateOutcome::ChannelClosed => return Ok(true), // pending still, WS disconnected
}
// `restart` calls process::exit and never returns — mark done first.
if tc.name == tn::RESTART {
info!(session_id = self.session_id, tool_call_id = tc.id, "restart approved (resume) — marking done then exiting");
chat_llm_tools::complete(pool, tc.id, "Riavvio avviato.", "string").await?;
em.tool_done(tc.id, "Riavvio avviato.".to_string(), "string".to_string()).await;
// Use _exit() to skip C atexit handlers (e.g. Metal GPU cleanup in
// whisper-rs/ggml, which aborts with SIGABRT and yields exit code 134
// instead of 255 — breaking the run.sh restart supervisor).
unsafe { libc::_exit(-1) }
}
// Re-run the persisted intent through the SAME dispatcher as a live turn
// (`execute_tool_call`), not the flat `build_execution`. This routes
// sub-agent tools (`execute_task` mode=sync, `execute_subtask`,
// `run_subtask`) through the recursive interception in `dispatch.rs`;
// `build_execution` alone does not know them and would fail with
// "Unknown tool: execute_task". Apply the RunContext working dir exactly
// like the live loop.
let effective_args = self.effective_args(&tc.name, &args).await;
let outcome = match self.execute_tool_call(
stack_id, config, tc.id, &tc.name, &effective_args, token, tx,
).await {
super::dispatch::DispatchResult::Outcome(o) => o,
// Clarification WS channel closed mid-resume — leave the tool pending
// so the next resume re-asks (mirrors the live turn's AbortPending).
super::dispatch::DispatchResult::AbortPending => return Ok(true),
};
// resume passes `None`: it does not accumulate ToolCallEvents nor re-emit
// FileChanged (only a live turn does). A /stop mid-resume returns Abort.
match self.record_tool_outcome(tc.id, &tc.name, &effective_args, outcome, &em, None).await? {
RecordFlow::Continue => {}
RecordFlow::Abort => return Ok(true),
}
}
Ok(true)
}
}
/// Shallowest stack depth that has more than one active (non-terminated) frame —
/// the top of an interrupted parallel sub-agent batch. Returns `None` for a linear
/// stack, where every depth has at most one active frame. Pure (see tests).
fn shallowest_parallel_depth(active: &[chat_sessions_stack::SessionStack]) -> Option<i64> {
let mut by_depth: std::collections::HashMap<i64, usize> = std::collections::HashMap::new();
for f in active {
*by_depth.entry(f.depth).or_default() += 1;
}
by_depth.iter()
.filter_map(|(depth, count)| (*count > 1).then_some(*depth))
.min()
}
#[cfg(test)]
mod tests {
use super::shallowest_parallel_depth;
use crate::db::chat_sessions_stack::SessionStack;
fn frame(id: i64, depth: i64, parent: Option<i64>) -> SessionStack {
SessionStack { id, agent_id: "agent".into(), depth, parent_tool_call_id: parent }
}
#[test]
fn linear_stack_is_not_a_batch() {
let frames = vec![frame(1, 0, None), frame(2, 1, Some(10)), frame(3, 2, Some(20))];
assert_eq!(shallowest_parallel_depth(&frames), None);
assert_eq!(shallowest_parallel_depth(&[]), None);
}
#[test]
fn detects_shallowest_multi_frame_depth() {
// Two siblings at depth 1 (parallel batch) plus a grandchild at depth 2.
let frames = vec![
frame(1, 0, None),
frame(2, 1, Some(10)), frame(3, 1, Some(11)),
frame(4, 2, Some(30)),
];
assert_eq!(shallowest_parallel_depth(&frames), Some(1));
}
#[test]
fn detects_deeper_batch_when_upper_levels_linear() {
let frames = vec![
frame(1, 0, None),
frame(2, 1, Some(10)),
frame(3, 2, Some(20)), frame(4, 2, Some(21)),
];
assert_eq!(shallowest_parallel_depth(&frames), Some(2));
}
}
+180
View File
@@ -0,0 +1,180 @@
use std::collections::HashMap;
use std::sync::Arc;
use sqlx::SqlitePool;
use tokio::sync::Mutex;
use crate::approval::ApprovalManager;
use crate::chat_event_bus::ChatEventBus;
use crate::clarification::ClarificationManager;
use crate::compactor::ContextCompactor;
use crate::config::DatetimeConfig;
use crate::db::{chat_sessions, chat_sessions_stack};
use crate::llm::LlmManager;
use crate::mcp::McpManager;
use crate::image_generate::ImageGeneratorManager;
use crate::memory::MemoryManager;
use crate::run_context::{RunContext, RunContextManager};
use crate::tool_discovery::ToolDiscovery;
use crate::tools::ToolRegistry;
use super::handler::ChatSessionHandler;
pub struct ChatSessionManager {
db: Arc<SqlitePool>,
llm_manager: Arc<LlmManager>,
max_history_messages: usize,
max_tool_rounds: usize,
max_parallel_subagents: usize,
max_tool_result_chars: Option<usize>,
datetime_config: DatetimeConfig,
tools: Arc<ToolRegistry>,
mcp: Arc<McpManager>,
approval: Arc<ApprovalManager>,
clarification: Arc<ClarificationManager>,
event_bus: Arc<ChatEventBus>,
memory_manager: Arc<MemoryManager>,
image_generator_manager: Arc<ImageGeneratorManager>,
/// Shared compactor instance, `None` when compaction is disabled.
compactor: Option<Arc<ContextCompactor>>,
run_context_manager: Arc<RunContextManager>,
/// Shared tool-discovery recorder, passed to every handler so each turn can
/// register the tools it actually offers to the LLM (see `ToolDiscovery`).
tool_discovery: Arc<ToolDiscovery>,
active: Mutex<HashMap<i64, Arc<ChatSessionHandler>>>,
}
impl ChatSessionManager {
pub fn new(
db: Arc<SqlitePool>,
llm_manager: Arc<LlmManager>,
max_history_messages: usize,
max_tool_rounds: usize,
max_parallel_subagents: usize,
max_tool_result_chars: Option<usize>,
datetime_config: DatetimeConfig,
tools: Arc<ToolRegistry>,
mcp: Arc<McpManager>,
approval: Arc<ApprovalManager>,
clarification: Arc<ClarificationManager>,
event_bus: Arc<ChatEventBus>,
memory_manager: Arc<MemoryManager>,
image_generator_manager: Arc<ImageGeneratorManager>,
compactor: Option<Arc<ContextCompactor>>,
run_context_manager: Arc<RunContextManager>,
tool_discovery: Arc<ToolDiscovery>,
) -> Self {
Self {
db,
llm_manager,
max_history_messages,
max_tool_rounds,
max_parallel_subagents,
max_tool_result_chars,
datetime_config,
tools,
mcp,
approval,
clarification,
event_bus,
memory_manager,
image_generator_manager,
compactor,
run_context_manager,
tool_discovery,
active: Mutex::new(HashMap::new()),
}
}
pub fn llm_manager(&self) -> Arc<LlmManager> {
Arc::clone(&self.llm_manager)
}
pub fn run_context_manager(&self) -> Arc<RunContextManager> {
Arc::clone(&self.run_context_manager)
}
/// Returns the live handler for `session_id` if it is currently loaded,
/// without creating a new one. Used by the API for in-place updates.
pub async fn active_handler(&self, session_id: i64) -> Option<Arc<ChatSessionHandler>> {
self.active.lock().await.get(&session_id).cloned()
}
pub async fn create_session(
&self,
agent_id: &str,
source: &str,
is_interactive: bool,
is_ephemeral: bool,
run_context: Option<&RunContext>,
) -> anyhow::Result<(i64, i64)> {
let session = chat_sessions::create(&self.db, agent_id, source, is_interactive, is_ephemeral).await?;
// Persist the RunContext at creation time so it is present before any handler
// is constructed (get_or_create_handler reads it once at construction).
if let Some(rc) = run_context {
chat_sessions::set_run_context(&self.db, session.id, Some(&rc.to_db())).await?;
}
let stack = chat_sessions_stack::create(
&self.db, session.id, "main", None, 0, None,
).await?;
Ok((session.id, stack.id))
}
/// Cancel the in-flight turn for `session_id` and clean up any pending
/// approvals and clarifications so their blocking awaits unblock immediately.
/// No-op if no handler is active for the session.
pub async fn cancel_session(&self, session_id: i64) {
let handler = self.active.lock().await.get(&session_id).cloned();
if let Some(h) = handler {
h.cancel();
h.cancel_pending_approvals().await;
h.cancel_pending_questions().await;
}
}
pub async fn get_or_create_handler(
&self,
session_id: i64,
) -> anyhow::Result<Arc<ChatSessionHandler>> {
{
let active = self.active.lock().await;
if let Some(h) = active.get(&session_id) {
return Ok(h.clone());
}
}
let session = chat_sessions::find_by_id(&self.db, session_id)
.await?
.ok_or_else(|| anyhow::anyhow!("session {session_id} not found"))?;
let run_context = session.run_context.as_deref().and_then(RunContext::from_db);
let handler = Arc::new(ChatSessionHandler::new(
session_id,
self.db.clone(),
Arc::clone(&self.llm_manager),
self.max_history_messages,
self.max_tool_rounds,
self.max_parallel_subagents,
self.max_tool_result_chars,
self.datetime_config.clone(),
session.agent_id,
session.source,
session.is_interactive,
session.is_ephemeral,
self.tools.clone(),
self.mcp.clone(),
Arc::clone(&self.approval),
Arc::clone(&self.clarification),
Arc::clone(&self.event_bus),
Arc::clone(&self.memory_manager),
Arc::clone(&self.image_generator_manager),
self.compactor.clone(),
run_context,
Arc::clone(&self.tool_discovery),
));
self.active.lock().await.insert(session_id, handler.clone());
Ok(handler)
}
}
+2
View File
@@ -0,0 +1,2 @@
pub mod handler;
pub mod manager;
+102
View File
@@ -0,0 +1,102 @@
//! The logical API surface of `Skald`: one accessor per manager, named after the
//! historical field, delegating into the domain bundle that now owns it. This is
//! the intentional surface consumers (frontend handlers, plugin context) use — the
//! bundles themselves stay internal.
//!
//! Now that the core is its own crate, this is a real boundary rather than a
//! convention: everything here is `pub` because the `skald` binary lives outside,
//! and everything not here is unreachable from it. Promote the block to a
//! `SkaldApi` trait if the shells ever need to mock it.
use std::sync::Arc;
use sqlx::SqlitePool;
use tokio::sync::RwLock;
use tokio_util::sync::CancellationToken;
use core_api::remote::RemoteAccess;
use core_api::system_bus::SystemEventBus;
use crate::approval::ApprovalManager;
use crate::chat_event_bus::ChatEventBus;
use crate::chat_hub::ChatHub;
use crate::clarification::ClarificationManager;
use crate::command::LlmCommandManager;
use crate::config_store::GlobalConfigManager;
use crate::cron::TaskManager;
use crate::elicitation::ElicitationManager;
use crate::image_generate::ImageGeneratorManager;
use crate::inbox::Inbox;
use crate::latex::LatexCompiler;
use crate::llm::LlmManager;
use crate::location::LocationManager;
use crate::mcp::McpManager;
use crate::memory::MemoryManager;
use crate::plugin::PluginManager;
use crate::projects::tickets::ProjectTicketManager;
use crate::projects::ProjectManager;
use crate::provider::ProviderRegistry;
use crate::run_context::RunContextManager;
use crate::secrets::SecretsStore;
use crate::session::manager::ChatSessionManager;
use crate::tic::TicManager;
use crate::tool_catalog::ToolCatalog;
use crate::tools::ToolRegistry;
use crate::transcribe::TranscribeManager;
use crate::tts::TtsManager;
use crate::users::UserManager;
use super::Skald;
impl Skald {
// Runtime / cross-cutting
pub fn db(&self) -> &Arc<SqlitePool> { &self.rt.db }
pub fn users(&self) -> &Arc<UserManager> { &self.rt.users }
pub fn config(&self) -> &Arc<GlobalConfigManager> { &self.rt.config }
pub fn config_properties(&self) -> &[core_api::ConfigSet] { &self.rt.config_properties }
pub fn system_bus(&self) -> &Arc<SystemEventBus> { &self.rt.system_bus }
pub fn event_bus(&self) -> &Arc<ChatEventBus> { &self.rt.event_bus }
pub fn shutdown_token(&self) -> &CancellationToken { &self.rt.shutdown_token }
// Models
pub fn provider_registry(&self) -> &Arc<ProviderRegistry> { &self.models.provider_registry }
pub fn llm_manager(&self) -> &Arc<LlmManager> { &self.models.llm_manager }
pub fn secrets(&self) -> &Arc<SecretsStore> { &self.models.secrets }
pub fn memory_manager(&self) -> &Arc<MemoryManager> { &self.models.memory_manager }
// Media
pub fn image_generator_manager(&self) -> &Arc<ImageGeneratorManager> { &self.media.image_generator_manager }
pub fn transcribe_manager(&self) -> &Arc<TranscribeManager> { &self.media.transcribe_manager }
pub fn tts_manager(&self) -> &Arc<TtsManager> { &self.media.tts_manager }
// Tools
pub fn tools(&self) -> &Arc<ToolRegistry> { &self.tools.tools }
pub fn catalog(&self) -> &ToolCatalog { &self.tools.catalog }
pub fn command_manager(&self) -> &Arc<LlmCommandManager> { &self.tools.command_manager }
// Integrations
pub fn mcp(&self) -> &Arc<McpManager> { &self.integrations.mcp }
pub fn plugin_manager(&self) -> &Arc<PluginManager> { &self.integrations.plugin_manager }
// Tasks
pub fn cron(&self) -> &Arc<TaskManager> { &self.tasks.cron }
pub fn projects(&self) -> &Arc<ProjectManager> { &self.tasks.projects }
pub fn ticket_manager(&self) -> &Arc<ProjectTicketManager> { &self.tasks.ticket_manager }
// Conversation
pub fn manager(&self) -> &Arc<ChatSessionManager> { &self.conversation.manager }
pub fn chat_hub(&self) -> &Arc<ChatHub> { &self.conversation.chat_hub }
pub fn run_context_manager(&self) -> &Arc<RunContextManager> { &self.conversation.run_context_manager }
pub fn tic_manager(&self) -> &Arc<TicManager> { &self.conversation.tic_manager }
// Interaction
pub fn approval(&self) -> &Arc<ApprovalManager> { &self.interaction.approval }
pub fn inbox(&self) -> &Inbox { &self.interaction.inbox }
pub fn clarification(&self) -> &Arc<ClarificationManager> { &self.interaction.clarification }
pub fn elicitation(&self) -> &Arc<ElicitationManager> { &self.interaction.elicitation }
// Infra
pub fn latex_compiler(&self) -> &LatexCompiler { &self.infra.latex_compiler }
pub fn location_manager(&self) -> &Arc<LocationManager> { &self.infra.location_manager }
pub fn remote(&self) -> &Arc<RwLock<Option<Arc<dyn RemoteAccess>>>> { &self.infra.remote }
}
+407
View File
@@ -0,0 +1,407 @@
//! Domain bundles: the managers, grouped by cohesion, that make up `Skald`.
//!
//! Each bundle owns a `build()` that constructs its managers (plus their startup
//! logging and non-fatal `seed_*` calls) from the shared [`Runtime`] and whatever
//! sibling bundles it depends on at construction time. Cross-bundle *cycles* are
//! not expressed here — they are resolved by the managers' `OnceLock` setters,
//! called in one place by [`super::wiring::wire`]. Bundle structs never hold
//! references to each other.
use std::sync::Arc;
use anyhow::Result;
use tracing::{debug, info, warn};
use core_api::remote::RemoteAccess;
use crate::approval::ApprovalManager;
use crate::chat_hub::ChatHub;
use crate::clarification::ClarificationManager;
use crate::command::LlmCommandManager;
use crate::compactor::ContextCompactor;
use crate::config::{CoreConfig, DatetimeConfig};
use crate::cron::TaskManager;
use crate::elicitation::ElicitationManager;
use crate::image_generate::ImageGeneratorManager;
use crate::inbox::Inbox;
use crate::latex::LatexCompiler;
use crate::llm::LlmManager;
use crate::location::LocationManager;
use crate::mcp::McpManager;
use crate::memory::MemoryManager;
use crate::plugin::PluginManager;
use crate::projects::tickets::ProjectTicketManager;
use crate::projects::ProjectManager;
use crate::provider::ProviderRegistry;
use crate::run_context::RunContextManager;
use crate::secrets::SecretsStore;
use crate::session::handler::{DEFAULT_MAX_PARALLEL_SUBAGENTS, DEFAULT_MAX_TOOL_ROUNDS};
use crate::session::manager::ChatSessionManager;
use crate::tic::TicManager;
use crate::tool_catalog::ToolCatalog;
use crate::tool_discovery::ToolDiscovery;
use crate::tools::ToolRegistry;
use crate::transcribe::TranscribeManager;
use crate::tts::TtsManager;
use tokio::sync::RwLock;
use core_api::plugin::Plugin;
use super::runtime::Runtime;
// ── Models: LLM/provider stack ──────────────────────────────────────────────
pub(super) struct Models {
pub(super) provider_registry: Arc<ProviderRegistry>,
pub(super) llm_manager: Arc<LlmManager>,
pub(super) secrets: Arc<SecretsStore>,
pub(super) memory_manager: Arc<MemoryManager>,
}
impl Models {
pub(super) async fn build(rt: &Runtime, config: &CoreConfig) -> Result<Self> {
let mut provider_registry = ProviderRegistry::new(Arc::clone(&rt.system_bus));
provider_registry.register_builtin(crate::llm::providers::openai::OpenAiProvider);
provider_registry.register_builtin(crate::llm::providers::anthropic::AnthropicProvider::new());
provider_registry.register_builtin(crate::llm::providers::openrouter::OpenRouterProvider::new());
provider_registry.register_builtin(crate::llm::providers::ollama::OllamaProvider::new());
provider_registry.register_builtin(crate::llm::providers::lm_studio::LmStudioProvider::new());
provider_registry.register_builtin(crate::llm::providers::deepseek::DeepSeekProvider::new());
provider_registry.register_builtin(crate::llm::providers::zai::ZaiProvider::new());
let provider_registry = Arc::new(provider_registry);
info!("provider registry ready ({} built-in providers)", provider_registry.all().len());
let log_flags = config.llm.requests_log.as_ref().filter(|r| r.enabled).map(|r| {
use crate::chatbot::logging::LogSaveFlags;
LogSaveFlags {
request_payload: r.request_payload_save,
response_payload: r.response_payload_save,
request_headers: r.request_header_save,
response_headers: r.response_header_save,
}
});
let llm_manager = LlmManager::new(Arc::clone(&rt.db), Arc::clone(&provider_registry), log_flags).await?;
let client_count = llm_manager.client_names().await.len().saturating_sub(1);
let default_client = llm_manager.default_name().await;
info!(clients = client_count, default = %default_client, "LLM clients loaded");
let secrets = SecretsStore::new(Arc::clone(&rt.db));
info!("secrets store ready");
let memory_manager = Arc::new(MemoryManager::new());
info!("memory manager ready");
Ok(Models { provider_registry, llm_manager, secrets, memory_manager })
}
}
// ── Media: transcription / TTS / image generation ───────────────────────────
pub(super) struct Media {
pub(super) image_generator_manager: Arc<ImageGeneratorManager>,
pub(super) transcribe_manager: Arc<TranscribeManager>,
pub(super) tts_manager: Arc<TtsManager>,
}
impl Media {
pub(super) async fn build(rt: &Runtime, models: &Models) -> Result<Self> {
let image_generator_manager = ImageGeneratorManager::new(
Arc::clone(&rt.db),
Arc::clone(&models.provider_registry),
"data",
).await?;
// Evaluate the await outside the `info!` macro: leaving the temporary
// `tracing::Value` from the field expression alive across the await
// makes the surrounding future non-Send, which Tauri's runtime rejects.
let image_generator_models = image_generator_manager.list_models_info().await.len();
info!(
db_backed = image_generator_models,
"image generator manager ready",
);
let transcribe_manager = TranscribeManager::new(
Arc::clone(&rt.db),
Arc::clone(&models.provider_registry),
Arc::clone(&rt.system_bus),
rt.shutdown_token.clone(),
).await?;
let transcribe_models = transcribe_manager.list_models_info().await.len();
info!(
db_backed = transcribe_models,
"transcribe manager ready",
);
let tts_manager = TtsManager::new(
Arc::clone(&rt.db),
Arc::clone(&models.provider_registry),
Arc::clone(&rt.system_bus),
rt.shutdown_token.clone(),
).await?;
let tts_models = tts_manager.list_models_info().await.len();
info!(
db_backed = tts_models,
"tts manager ready",
);
Ok(Media { image_generator_manager, transcribe_manager, tts_manager })
}
}
// ── Integrations: MCP + plugins ─────────────────────────────────────────────
pub(super) struct Integrations {
pub(super) mcp: Arc<McpManager>,
pub(super) plugin_manager: Arc<PluginManager>,
}
impl Integrations {
/// Builds the MCP manager (its `initialize()` is deferred to `spawn_background`,
/// after the elicitation handler is wired) and the plugin manager (plugins are
/// injected by `main.rs`; `start_enabled()` runs later, from `WebFrontend`).
pub(super) fn build(rt: &Runtime, plugins: Vec<Arc<dyn Plugin>>) -> Self {
let mcp = Arc::new(McpManager::new(Arc::clone(&rt.db), rt.shutdown_token.clone(), "data"));
let mut plugin_manager = PluginManager::new(Arc::clone(&rt.db));
for plugin in plugins {
plugin_manager.register_arc(plugin);
}
info!("plugins registered");
let plugin_manager = Arc::new(plugin_manager);
Integrations { mcp, plugin_manager }
}
}
// ── Tasks: cron + projects/tickets ──────────────────────────────────────────
pub(super) struct Tasks {
pub(super) cron: Arc<TaskManager>,
pub(super) projects: Arc<ProjectManager>,
pub(super) ticket_manager: Arc<ProjectTicketManager>,
}
impl Tasks {
/// Built before `Tools` so cron tools can capture the `TaskManager`.
pub(super) fn build(rt: &Runtime, config: &CoreConfig) -> Self {
let cron_tz = config.timezone.as_deref().and_then(|s| {
match s.parse::<chrono_tz::Tz>() {
Ok(tz) => { info!("timezone: using {s}"); Some(tz) }
Err(_) => { warn!("timezone: unknown value '{s}', falling back to local time"); None }
}
});
let cron = TaskManager::new(Arc::clone(&rt.db), cron_tz, Arc::clone(&rt.system_bus));
let ticket_manager = ProjectTicketManager::new(Arc::clone(&rt.db));
let projects = Arc::new(ProjectManager::new(Arc::clone(&rt.db)));
info!("project manager ready");
Tasks { cron, projects, ticket_manager }
}
}
// ── Tools: registry + catalog + slash commands ──────────────────────────────
pub(super) struct Tools {
pub(super) tools: Arc<ToolRegistry>,
pub(super) catalog: ToolCatalog,
pub(super) command_manager: Arc<LlmCommandManager>,
}
impl Tools {
/// Captures sibling managers (mcp, plugins, cron, secrets) into the tool
/// registry. `execute_task` is deliberately NOT registered here — it is injected
/// per interactive session by `ChatHub::send_message`.
pub(super) fn build(integrations: &Integrations, tasks: &Tasks, models: &Models) -> Self {
let mut tool_registry = ToolRegistry::new();
crate::tools::fs::register_all(&mut tool_registry);
tool_registry.register(crate::tools::ast_outline::AstOutline::new());
tool_registry.register(crate::tools::exec::ExecuteCmd);
tool_registry.register(crate::tools::read_notification::ReadNotification);
tool_registry.register(crate::tools::restart::Restart);
// Unified listing / toggling across mcp, plugins, cron (+ agents for list).
tool_registry.register(crate::tools::list_items::ListItems::new(
Arc::clone(&integrations.mcp), Arc::clone(&integrations.plugin_manager), Arc::clone(&tasks.cron)));
tool_registry.register(crate::tools::toggle_item::ToggleItem::new(
Arc::clone(&integrations.mcp), Arc::clone(&integrations.plugin_manager), Arc::clone(&tasks.cron)));
tool_registry.register(crate::tools::register_mcp::RegisterMcp::new(Arc::clone(&integrations.mcp)));
tool_registry.register(crate::tools::register_mcp::DeleteMcp::new(Arc::clone(&integrations.mcp)));
tool_registry.register(crate::tools::cron_jobs::DeleteCronJob(Arc::clone(&tasks.cron)));
tool_registry.register(crate::tools::set_secret::SetSecret(Arc::clone(&models.secrets)));
tool_registry.register(crate::tools::list_secrets::ListSecrets(Arc::clone(&models.secrets)));
tool_registry.register(crate::tools::configure_plugin::ConfigurePlugin(Arc::clone(&integrations.plugin_manager)));
// Tools contributed by plugins (plugin.md §11), via `Plugin::tools()`.
// The core never names a plugin crate: each one hands over whatever tools
// it wants, bound to its own handle. They are built before the plugins'
// runloops start, so they must tolerate being called while stopped.
for plugin in integrations.plugin_manager.all() {
let id = plugin.id().to_string();
let tools = Arc::clone(plugin).tools();
if tools.is_empty() {
continue;
}
let n = tools.len();
for tool in tools {
tool_registry.register_arc(tool);
}
info!(plugin = %id, count = n, "plugin tools registered");
}
debug!("tool registry built");
let tools = Arc::new(tool_registry);
let catalog = ToolCatalog::new(Arc::clone(&tools), Arc::clone(&integrations.mcp));
let command_manager = Arc::new(LlmCommandManager::new());
Tools { tools, catalog, command_manager }
}
}
// ── Interaction: approval + inbox + clarification + elicitation ─────────────
pub(super) struct Interaction {
pub(super) approval: Arc<ApprovalManager>,
pub(super) inbox: Inbox,
pub(super) clarification: Arc<ClarificationManager>,
pub(super) elicitation: Arc<ElicitationManager>,
}
impl Interaction {
pub(super) async fn build(rt: &Runtime, tools: &Tools) -> Result<Self> {
let approval = Arc::new(ApprovalManager::new(Arc::clone(&rt.db), rt.global_tx.clone()));
if let Err(e) = approval.seed_defaults().await {
warn!(error = %e, "failed to seed default approval rules (non-fatal)");
}
if let Err(e) = approval.migrate_legacy_fs_rules().await {
warn!(error = %e, "failed to migrate legacy filesystem rules (non-fatal)");
}
if let Err(e) = approval.seed_fs_path_rules().await {
warn!(error = %e, "failed to seed File System path rules (non-fatal)");
}
if let Err(e) = approval.seed_default_catch_all().await {
warn!(error = %e, "failed to seed default catch-all rule (non-fatal)");
}
info!("approval manager ready");
let clarification = ClarificationManager::new(rt.global_tx.clone());
let elicitation = ElicitationManager::new(rt.global_tx.clone());
let inbox = Inbox::new(
Arc::clone(&approval),
Arc::clone(&clarification),
Arc::clone(&elicitation),
Arc::clone(&tools.tools),
);
Ok(Interaction { approval, inbox, clarification, elicitation })
}
}
// ── Conversation: session manager + chat hub + run context + TIC ────────────
pub(super) struct Conversation {
pub(super) manager: Arc<ChatSessionManager>,
pub(super) chat_hub: Arc<ChatHub>,
pub(super) run_context_manager: Arc<RunContextManager>,
/// TIC lives here (rather than in `Tasks`) because it is constructed from and
/// drives the conversation stack (session manager + chat hub + run context);
/// this keeps every bundle a single-shot `build()` with no two-phase init.
pub(super) tic_manager: Arc<TicManager>,
}
impl Conversation {
#[allow(clippy::too_many_arguments)]
pub(super) async fn build(
rt: &Runtime,
models: &Models,
media: &Media,
tools: &Tools,
integrations: &Integrations,
interaction: &Interaction,
config: &CoreConfig,
) -> Result<Self> {
let run_context_manager =
Arc::new(RunContextManager::new(Arc::clone(&rt.db), Arc::clone(&interaction.approval)));
if let Err(e) = run_context_manager.seed_defaults().await {
warn!(error = %e, "failed to seed default permission group (non-fatal)");
}
info!("run_context manager ready");
let compactor = config.llm.compaction.as_ref().map(|cfg| {
info!(
threshold_tokens = cfg.threshold_tokens,
keep_recent = cfg.keep_recent,
?cfg.strength,
"context compactor enabled"
);
Arc::new(ContextCompactor::new(
cfg.clone(),
Arc::clone(&models.llm_manager),
Arc::clone(&rt.event_bus),
))
});
if compactor.is_none() {
info!("context compactor disabled (no compaction config)");
}
let manager = Arc::new(ChatSessionManager::new(
Arc::clone(&rt.db),
Arc::clone(&models.llm_manager),
config.llm.max_history_messages,
config.llm.max_tool_rounds.unwrap_or(DEFAULT_MAX_TOOL_ROUNDS),
config.llm.max_parallel_subagents.unwrap_or(DEFAULT_MAX_PARALLEL_SUBAGENTS),
config.llm.max_tool_result_chars,
DatetimeConfig { timezone: config.timezone.clone(), ..config.llm.datetime },
Arc::clone(&tools.tools),
Arc::clone(&integrations.mcp),
Arc::clone(&interaction.approval),
Arc::clone(&interaction.clarification),
Arc::clone(&rt.event_bus),
Arc::clone(&models.memory_manager),
Arc::clone(&media.image_generator_manager),
compactor,
Arc::clone(&run_context_manager),
Arc::new(ToolDiscovery::new(Arc::clone(&rt.db))),
));
let chat_hub = ChatHub::new(
Arc::clone(&rt.db),
Arc::clone(&manager),
Arc::clone(&interaction.approval),
rt.global_tx.clone(),
rt.shutdown_token.clone(),
);
chat_hub.register("web").await;
chat_hub.register("talk").await;
let tic_manager = TicManager::new(
Arc::clone(&rt.db),
Arc::clone(&manager),
Arc::clone(&chat_hub),
config.tic.clone(),
Arc::clone(&rt.config),
Arc::clone(&run_context_manager),
Arc::clone(&rt.system_bus),
);
Ok(Conversation { manager, chat_hub, run_context_manager, tic_manager })
}
}
// ── Infra: leftover singletons ──────────────────────────────────────────────
pub(super) struct Infra {
pub(super) latex_compiler: LatexCompiler,
pub(super) location_manager: Arc<LocationManager>,
pub(super) remote: Arc<RwLock<Option<Arc<dyn RemoteAccess>>>>,
}
impl Infra {
pub(super) fn build() -> Self {
Infra {
latex_compiler: LatexCompiler::new(),
location_manager: Arc::new(LocationManager::new()),
remote: Arc::new(RwLock::new(None)),
}
}
}
+99
View File
@@ -0,0 +1,99 @@
//! `Skald` — the headless application core.
//!
//! `Skald` owns every manager but is no longer a God Object: the ~30 managers are
//! grouped into a cross-cutting [`Runtime`] context plus eight cohesive domain
//! bundles (see [`bundles`]). Construction is a staged composition root (each bundle
//! has its own `build()`); the construction cycles are resolved in one place by
//! [`wiring::wire`]; every background task is registered with a [`TaskSupervisor`]
//! so shutdown joins them uniformly. The frontend and plugin context consume `Skald`
//! only through the accessor methods in [`accessors`], never its fields — that
//! accessor surface is the logical boundary a future `skald-core` crate would keep.
use std::sync::Arc;
use anyhow::Result;
use sqlx::SqlitePool;
use tracing::info;
use core_api::plugin::Plugin;
use super::config::CoreConfig;
mod accessors;
mod bundles;
mod runtime;
mod supervisor;
mod wiring;
use bundles::{Conversation, Infra, Integrations, Interaction, Media, Models, Tasks, Tools};
use runtime::Runtime;
use wiring::{spawn_background, wire};
pub struct Skald {
rt: Runtime,
models: Models,
media: Media,
tools: Tools,
integrations: Integrations,
tasks: Tasks,
conversation: Conversation,
interaction: Interaction,
infra: Infra,
}
impl Skald {
pub async fn new(pool: Arc<SqlitePool>, config: &CoreConfig, plugins: Vec<Arc<dyn Plugin>>) -> Result<Arc<Self>> {
let discovered = super::agents::discover()?;
info!(
count = discovered.len(),
agents = discovered.iter().map(|a| a.id.as_str()).collect::<Vec<_>>().join(", "),
"agents discovered"
);
// ── Composition root: build the runtime context, then each domain bundle
// in dependency order. `Tasks` precedes `Tools` (tools capture cron);
// `Interaction` and `Conversation` come last (they need the tool registry
// and each other's managers).
let rt = Runtime::bootstrap(pool);
let models = Models::build(&rt, config).await?;
let media = Media::build(&rt, &models).await?;
let integrations = Integrations::build(&rt, plugins);
let tasks = Tasks::build(&rt, config);
let tools = Tools::build(&integrations, &tasks, &models);
let interaction = Interaction::build(&rt, &tools).await?;
let conversation = Conversation::build(&rt, &models, &media, &tools, &integrations, &interaction, config).await?;
let infra = Infra::build();
// Resolve construction cycles, then start background tasks.
wire(&tasks, &conversation, &integrations, &interaction);
spawn_background(&rt, &tasks, &conversation, &integrations, config);
let skald = Arc::new(Skald {
rt, models, media, tools, integrations, tasks, conversation, interaction, infra,
});
// Inject the fully-constructed instance into the plugin manager — the one
// Arc<Skald> back-reference. start_enabled()/start_config_watcher() run later,
// from WebFrontend::start, once the router factory is wired.
skald.plugin_manager().set_skald(Arc::clone(&skald));
Ok(skald)
}
pub fn subscribe_chat_events(&self) -> tokio::sync::broadcast::Receiver<core_api::bus::BusEvent> {
self.rt.event_bus.subscribe()
}
pub fn subscribe_system_events(&self) -> tokio::sync::broadcast::Receiver<core_api::system_bus::SystemEvent> {
self.rt.system_bus.subscribe()
}
pub async fn shutdown(self: Arc<Self>) {
self.rt.shutdown_token.cancel();
self.rt.supervisor.join_all(tokio::time::Duration::from_secs(10)).await;
self.integrations.plugin_manager.stop_all().await;
// Last: every user key leaves RAM. A restarted box is opaque again until
// each user unlocks their own database (§9).
self.rt.users.lock_all().await;
}
}
+70
View File
@@ -0,0 +1,70 @@
//! Cross-cutting runtime context.
//!
//! `Runtime` holds the primitives every domain bundle needs: the DB pool, the
//! global config manager, the two event buses, the server→client broadcast
//! channel, the shutdown token and the background-task supervisor. It is built
//! first and passed by reference into each bundle builder, so no bundle has to
//! depend on another purely to reach a shared primitive. This is also the natural
//! seam a future extracted `skald-core` crate would expose as its root context.
use std::sync::Arc;
use sqlx::SqlitePool;
use tokio::sync::broadcast;
use tokio_util::sync::CancellationToken;
use tracing::info;
use core_api::events::GlobalEvent;
use core_api::system_bus::SystemEventBus;
use crate::chat_event_bus::ChatEventBus;
use crate::config_store::GlobalConfigManager;
use crate::users::UserManager;
use super::supervisor::TaskSupervisor;
pub(super) struct Runtime {
/// The registry pool (`system.db`). Still the only pool anything reads or
/// writes: nothing has moved to per-user pools yet. `users` owns those.
pub(super) db: Arc<SqlitePool>,
pub(super) users: Arc<UserManager>,
pub(super) config: Arc<GlobalConfigManager>,
pub(super) config_properties: Vec<core_api::ConfigSet>,
pub(super) system_bus: Arc<SystemEventBus>,
pub(super) event_bus: Arc<ChatEventBus>,
/// Server→client push channel (`ServerEvent` wrapped in `GlobalEvent`). Shared
/// into approval / clarification / elicitation / chat_hub; consumed by the WS
/// handlers. Hoisted here so it exists before the bundles that need it.
pub(super) global_tx: broadcast::Sender<GlobalEvent>,
pub(super) shutdown_token: CancellationToken,
pub(super) supervisor: Arc<TaskSupervisor>,
}
impl Runtime {
/// Wires the cross-cutting primitives. Infallible.
pub(super) fn bootstrap(pool: Arc<SqlitePool>) -> Self {
let config = Arc::new(GlobalConfigManager::new(Arc::clone(&pool)));
let users = Arc::new(UserManager::new(Arc::clone(&pool)));
let system_bus = Arc::new(SystemEventBus::new());
info!("system event bus ready");
let event_bus = Arc::new(ChatEventBus::new());
info!("chat event bus ready");
let (global_tx, _) = broadcast::channel::<GlobalEvent>(512);
Runtime {
db: pool,
users,
config,
config_properties: vec![crate::tic::config_set()],
system_bus,
event_bus,
global_tx,
shutdown_token: CancellationToken::new(),
supervisor: TaskSupervisor::new(),
}
}
}
+59
View File
@@ -0,0 +1,59 @@
//! Background-task supervision.
//!
//! Every long-lived task spawned during `Skald::new` is registered here by name so
//! that `Skald::shutdown` can join them all against a single deadline and report any
//! laggards individually. This replaces the previous `bg_handles` vec, which only
//! tracked a subset of the spawned tasks (leaving the log-cleanup loop and
//! `mcp.initialize` fire-and-forget and never awaited).
use std::future::Future;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::task::JoinHandle;
use tracing::warn;
/// Tracks named background-task handles for graceful shutdown.
pub struct TaskSupervisor {
handles: Mutex<Vec<(&'static str, JoinHandle<()>)>>,
}
impl TaskSupervisor {
pub fn new() -> Arc<Self> {
Arc::new(Self { handles: Mutex::new(Vec::new()) })
}
/// Spawn a named future and track its handle.
pub fn spawn<F>(&self, name: &'static str, fut: F)
where
F: Future<Output = ()> + Send + 'static,
{
self.handles.lock().unwrap().push((name, tokio::spawn(fut)));
}
/// Adopt an already-spawned handle (for managers whose `start()` returns one).
pub fn adopt_one(&self, name: &'static str, handle: JoinHandle<()>) {
self.handles.lock().unwrap().push((name, handle));
}
/// Adopt a batch of handles (e.g. `cron.start()` returns `Vec<JoinHandle<()>>`).
pub fn adopt(&self, name: &'static str, handles: Vec<JoinHandle<()>>) {
let mut guard = self.handles.lock().unwrap();
for h in handles {
guard.push((name, h));
}
}
/// Join all tracked tasks against a shared deadline, logging any that do not
/// finish in time by name. Dropping a timed-out `JoinHandle` does not abort the
/// task; every task is already signalled via the shutdown `CancellationToken`.
pub async fn join_all(&self, timeout: Duration) {
let handles = std::mem::take(&mut *self.handles.lock().unwrap());
let deadline = tokio::time::Instant::now() + timeout;
for (name, handle) in handles {
if tokio::time::timeout_at(deadline, handle).await.is_err() {
warn!(task = name, "background task did not finish within shutdown deadline");
}
}
}
}
+101
View File
@@ -0,0 +1,101 @@
//! Post-construction wiring: the `OnceLock` cycle-breakers and the background-task
//! spawns, each concentrated in one readable place instead of being scattered
//! through the constructor.
use std::sync::Arc;
use tracing::info;
use crate::config::CoreConfig;
use crate::elicitation::ElicitationBridge;
use super::bundles::{Conversation, Integrations, Interaction, Tasks};
use super::runtime::Runtime;
/// Resolves the construction cycles (`cron ↔ session ↔ hub`, `ticket → cron`,
/// `mcp → elicitation`) via the managers' `OnceLock` setters.
pub(super) fn wire(
tasks: &Tasks,
conversation: &Conversation,
integrations: &Integrations,
interaction: &Interaction,
) {
tasks.cron.set_session(Arc::clone(&conversation.manager));
tasks.cron.set_hub(Arc::clone(&conversation.chat_hub));
tasks.cron.set_self_arc(Arc::clone(&tasks.cron));
tasks.ticket_manager.set_task_manager(Arc::clone(&tasks.cron));
conversation.chat_hub.set_task_mgr(Arc::clone(&tasks.cron));
integrations.mcp.set_elicitation_handler(ElicitationBridge::new(Arc::clone(&interaction.elicitation)));
info!("ChatHub initialised");
}
/// Spawns every long-lived background task, each registered by name with the
/// supervisor so it is joined on shutdown. MCP `initialize()` is spawned here —
/// after `wire()` has installed the elicitation handler — so stdio servers start
/// with a handler for server-initiated `elicitation/create` requests.
pub(super) fn spawn_background(
rt: &Runtime,
tasks: &Tasks,
conversation: &Conversation,
integrations: &Integrations,
config: &CoreConfig,
) {
// LLM request-log retention/cleanup — first run 1 min after startup, then 12h.
if let Some(cfg) = config.llm.requests_log.clone().filter(|r| r.enabled) {
rt.supervisor.adopt_one(
"llm-log-cleanup",
crate::db::llm_requests::cleanup::spawn(
Arc::clone(&rt.db),
cfg,
rt.shutdown_token.clone(),
),
);
}
// Session-cancellation subscriber: fans SessionCancelled events on the system
// bus into cancel_session() so any in-flight turn / approval / clarification
// all unblock.
{
let manager_ref = Arc::clone(&conversation.manager);
let mut rx = rt.system_bus.subscribe();
let sd = rt.shutdown_token.clone();
rt.supervisor.spawn("session-cancel", async move {
loop {
tokio::select! {
_ = sd.cancelled() => break,
event = rx.recv() => match event {
Ok(core_api::system_bus::SystemEvent::SessionCancelled { session_id }) => {
manager_ref.cancel_session(session_id).await;
}
Ok(_) => {}
Err(_) => break,
}
}
}
});
}
// MCP servers connect in the background. `initialize()` does not itself observe
// the cancellation token, so race it against shutdown: on cancel the task exits
// promptly (dropping the in-flight connection attempts) instead of blocking the
// shutdown join until the deadline.
{
let mcp = Arc::clone(&integrations.mcp);
let sd = rt.shutdown_token.clone();
rt.supervisor.spawn("mcp-init", async move {
tokio::select! {
_ = sd.cancelled() => {}
_ = mcp.initialize() => {}
}
});
}
rt.supervisor.adopt("cron", Arc::clone(&tasks.cron).start(rt.shutdown_token.clone()));
info!("cron scheduler started");
rt.supervisor.adopt_one(
"ticket-listener",
Arc::clone(&tasks.ticket_manager).start_listener(Arc::clone(&rt.system_bus), rt.shutdown_token.clone()),
);
rt.supervisor.adopt_one("tic", Arc::clone(&conversation.tic_manager).start(rt.shutdown_token.clone()));
info!("TicManager started");
}
+253
View File
@@ -0,0 +1,253 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use sqlx::SqlitePool;
use tokio::sync::mpsc;
use tracing::{info, warn};
use core_api::{ConfigProperty, ConfigSet, PropertyType};
use core_api::system_bus::{SystemEvent, SystemEventBus};
use crate::chat_hub::ChatHub;
use crate::config::TicConfig;
use crate::config_store::GlobalConfigManager;
use crate::db::mcp_events;
use crate::run_context::{RunContext, RunContextManager};
use crate::session::manager::ChatSessionManager;
const TIC_SOURCE: &str = "tic";
const TIC_AGENT: &str = "tic";
pub const TIC_ENABLED_KEY: &str = "tic.enabled";
pub const TIC_SECURITY_GROUP_KEY: &str = "tic.security_group";
pub const TIC_INTERVAL_MINUTES_KEY: &str = "tic.interval_minutes";
pub fn config_set() -> ConfigSet {
ConfigSet {
name: "TIC Agent".into(),
description: "TIC is a background agent that monitors all async events generated by connected MCP servers (new emails, calendar updates, WhatsApp messages, etc.). It reads your notification rules from data/notifications.md and your memory to decide — via an LLM call — which events are worth surfacing. Relevant notifications are forwarded to the home agent set via /sethome.".into(),
properties: vec![
ConfigProperty {
key: TIC_ENABLED_KEY.into(),
name: "Enabled".into(),
description: "Enable or disable the TIC agent. When disabled, no MCP events are processed.".into(),
property_type: PropertyType::Bool,
default_value: Some("true".into()),
},
ConfigProperty {
key: TIC_SECURITY_GROUP_KEY.into(),
name: "Security Group".into(),
description: "Tool permission group applied to each TIC agent session. Leave empty to use the default group.".into(),
property_type: PropertyType::SecurityGroup,
default_value: None,
},
ConfigProperty {
key: TIC_INTERVAL_MINUTES_KEY.into(),
name: "Check Interval (minutes)".into(),
description: "How often TIC runs, in minutes. Leave empty to use the value from config.yml (tic.interval_secs).".into(),
property_type: PropertyType::Int,
default_value: Some("15".into()),
},
],
}
}
pub struct TicManager {
db: Arc<SqlitePool>,
session_mgr: Arc<ChatSessionManager>,
hub: Arc<ChatHub>,
config: TicConfig,
config_store: Arc<GlobalConfigManager>,
run_context_manager: Arc<RunContextManager>,
system_bus: Arc<SystemEventBus>,
/// Guards against concurrent ticks (e.g. if a tick takes longer than the interval).
running: AtomicBool,
}
impl TicManager {
pub fn new(
db: Arc<SqlitePool>,
session_mgr: Arc<ChatSessionManager>,
hub: Arc<ChatHub>,
config: TicConfig,
config_store: Arc<GlobalConfigManager>,
run_context_manager: Arc<RunContextManager>,
system_bus: Arc<SystemEventBus>,
) -> Arc<Self> {
Arc::new(Self {
db,
session_mgr,
hub,
config,
config_store,
run_context_manager,
system_bus,
running: AtomicBool::new(false),
})
}
/// Force a tick immediately, ignoring the running guard.
/// Intended for manual triggering (e.g. via the `/api/tic/trigger` endpoint).
pub async fn tick_now(self: Arc<Self>) {
if let Err(e) = self.run_tick().await {
warn!(error = %e, "TicManager: forced tick failed");
}
}
/// Spawn the background timer.
/// Subscribes to ConfigKeyUpdated so the interval can be changed at runtime.
pub fn start(self: Arc<Self>, shutdown: tokio_util::sync::CancellationToken) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
let mut interval_secs = self.effective_interval_secs().await;
info!("TicManager started (interval={}s, batch={})", interval_secs, self.config.batch_size);
let mut timer = tokio::time::interval(Duration::from_secs(interval_secs));
timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut sys_rx = self.system_bus.subscribe();
loop {
tokio::select! {
_ = shutdown.cancelled() => {
info!("TicManager: stopping");
break;
}
res = sys_rx.recv() => {
if let Ok(SystemEvent::ConfigKeyUpdated { key, new_value, .. }) = res {
if key == TIC_INTERVAL_MINUTES_KEY {
if let Ok(mins) = new_value.parse::<u64>() {
let new_secs = mins.max(1) * 60;
if new_secs != interval_secs {
interval_secs = new_secs;
timer = tokio::time::interval(Duration::from_secs(interval_secs));
timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
info!(secs = interval_secs, "TicManager: interval updated");
}
}
}
}
}
_ = timer.tick() => {
self.tick().await;
}
}
}
})
}
async fn effective_interval_secs(&self) -> u64 {
if let Ok(Some(val)) = self.config_store.get(TIC_INTERVAL_MINUTES_KEY).await {
if let Ok(mins) = val.parse::<u64>() {
if mins > 0 {
return mins * 60;
}
}
}
self.config.interval_secs
}
async fn is_enabled(&self) -> bool {
match self.config_store.get(TIC_ENABLED_KEY).await {
Ok(Some(v)) => v != "false",
_ => true,
}
}
async fn tick(&self) {
if !self.is_enabled().await {
return;
}
// Prevent concurrent ticks.
if self.running.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst).is_err() {
warn!("TicManager: previous tick still running, skipping");
return;
}
let result = self.run_tick().await;
self.running.store(false, Ordering::SeqCst);
if let Err(e) = result {
warn!(error = %e, "TicManager: tick failed");
}
}
async fn run_tick(&self) -> anyhow::Result<()> {
// 1. Fetch the oldest N unprocessed events.
let events = mcp_events::pending_limited(&self.db, self.config.batch_size).await?;
if events.is_empty() {
return Ok(());
}
info!(count = events.len(), "TicManager: processing event batch");
// 2. Mark as processed BEFORE running the agent — avoids double-processing
// if the process crashes mid-turn.
let ids: Vec<i64> = events.iter().map(|e| e.id).collect();
mcp_events::mark_processed(&self.db, &ids).await?;
// 3. Serialize events into the agent prompt.
let prompt = build_prompt(&events);
// 4. Create a fresh ephemeral session (agent_id = "tic", source = "tic").
// We bypass ChatHub entirely — TIC is not a user-facing source and should
// not appear in the sources table or consume a broadcast channel.
let (session_id, _) = self.session_mgr.create_session(TIC_AGENT, TIC_SOURCE, false, true, None).await?;
let handler = self.session_mgr.get_or_create_handler(session_id).await?;
handler.set_auto_deny_approvals();
// 5. Apply run context if configured in DB.
if let Ok(Some(rc_id)) = self.config_store.get(TIC_SECURITY_GROUP_KEY).await {
if !rc_id.is_empty() {
let rc = RunContext::with_security_group(Some(rc_id.clone()));
if let Err(e) = self.run_context_manager.set_session_run_context(session_id, Some(&rc)).await {
warn!(error = %e, rc_id, "TicManager: failed to set run context");
}
}
}
// 6. Sink for session events — nobody subscribes; drop the receiver immediately
// so the channel is drained without buffering.
let (tx, _rx) = mpsc::channel(32);
let notify = crate::tools::notify::make_tool(Arc::clone(&self.hub), "TIC");
handler.handle_message(&prompt, None, None, None, None, vec![notify], std::collections::HashMap::new(), tx, true, None, None).await?;
info!(session_id, count = events.len(), "TicManager: tick complete");
Ok(())
}
}
// ── Prompt builder ─────────────────────────────────────────────────────────────
fn build_prompt(events: &[crate::db::mcp_events::McpEvent]) -> String {
use std::fmt::Write;
let n = events.len();
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S UTC");
let mut out = format!("[TIC] {n} pending event(s) — {now}\n");
for (i, ev) in events.iter().enumerate() {
let _ = write!(
out,
"\n=== Event {}/{n} ===\nSource: {}\nType: {}\nReceived: {}\nPayload:\n{}\n",
i + 1,
ev.source,
ev.method,
ev.created_at,
indent_payload(&ev.payload),
);
}
out
}
/// Pretty-print a JSON payload with 2-space indent, falling back to raw string.
fn indent_payload(payload: &str) -> String {
if let Ok(v) = serde_json::from_str::<serde_json::Value>(payload) {
if let Ok(pretty) = serde_json::to_string_pretty(&v) {
return pretty.lines().map(|l| format!(" {l}")).collect::<Vec<_>>().join("\n");
}
}
format!(" {payload}")
}
+119
View File
@@ -0,0 +1,119 @@
use std::collections::HashMap;
use std::sync::Arc;
use serde::Serialize;
use serde_json::Value;
use crate::mcp::McpManager;
use crate::tools::{ToolCategory, ToolDescriptionLength, ToolRegistry};
use crate::tools::tool_names as tn;
#[derive(Debug, Clone, Serialize)]
pub struct ToolInfo {
pub name: String,
pub description: String,
pub source: String,
pub server: Option<String>,
pub category: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct McpServerMeta {
pub friendly_name: Option<String>,
pub description: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct AllTools {
pub built_in: Vec<ToolInfo>,
pub mcp: Vec<ToolInfo>,
/// server internal name → metadata (friendly_name, description).
/// Populated by the API handler via a DB query; empty when constructed here.
#[serde(default)]
pub mcp_servers: HashMap<String, McpServerMeta>,
}
pub struct ToolCatalog {
tools: Arc<ToolRegistry>,
mcp: Arc<McpManager>,
}
impl ToolCatalog {
pub fn new(tools: Arc<ToolRegistry>, mcp: Arc<McpManager>) -> Self {
Self { tools, mcp }
}
pub fn list_all(&self) -> AllTools {
let mut built_in: Vec<ToolInfo> = self.tools
.list_all()
.into_iter()
.map(|(name, description)| {
let category = self.tools.category_of(&name).map(category_str);
ToolInfo { name, description, source: "built-in".into(), server: None, category }
})
.collect();
for (name, description, category) in Self::synthetic_tools() {
built_in.push(ToolInfo {
name: (*name).to_string(),
description: (*description).to_string(),
source: "built-in".into(),
server: None,
category: Some((*category).to_string()),
});
}
built_in.sort_by(|a, b| a.name.cmp(&b.name));
let mcp: Vec<ToolInfo> = self.mcp
.tools()
.into_iter()
.map(|t| ToolInfo {
name: t.tool_id(),
description: t.description,
source: "mcp".into(),
server: Some(t.server_name),
category: None,
})
.collect();
AllTools { built_in, mcp, mcp_servers: HashMap::new() }
}
pub fn describe_call(&self, name: &str, args: &Value, length: ToolDescriptionLength) -> String {
self.tools.describe_call(name, args, length)
}
/// Core-owned tools that are injected per-session outside the `ToolRegistry`
/// (interface tools + the provider-gated `image_generate`), listed statically
/// so they can be pre-configured in the Security-groups UI *before* first use.
///
/// This is a best-effort eager list — correctness does not depend on it being
/// complete: `ToolDiscovery` surfaces any tool that is actually offered, and
/// the catch-all `* require` gates anything not yet configured. Only names the
/// core legitimately owns belong here; plugin/provider tool names are left to
/// discovery so core stays decoupled from them.
fn synthetic_tools() -> &'static [(&'static str, &'static str, &'static str)] {
&[
(tn::EXECUTE_TASK, "Delegate to / schedule a sub-agent (cron, sync=inline sub-agent, async=background).", "subagent"),
(tn::EXECUTE_SUBTASK, "Run a synchronous sub-task inside a background session.", "subagent"),
(tn::UPDATE_SCRATCHPAD, "Write a key-value note into the session scratchpad.", "introspection"),
(tn::ASK_USER_CLARIFICATION, "Pause and ask the user a clarification question.", "introspection"),
(tn::WRITE_TODOS, "Record and update the agent's private per-turn task list.", "introspection"),
(tn::ACTIVATE_TOOLS, "Unlock an MCP server's tools or the built-in config tool group for this session.", "config"),
(tn::NOTIFY, "Send a proactive notification to the user (background/event-triage sessions).", "introspection"),
(tn::SHOW_FILE_TO_USER, "Open a file in the user's viewer (web/mobile sessions).", "introspection"),
(tn::IMAGE_GENERATE, "Generate an image from a text prompt (requires an image provider).", "config"),
]
}
}
fn category_str(cat: ToolCategory) -> String {
match cat {
ToolCategory::Filesystem => "filesystem",
ToolCategory::Shell => "shell",
ToolCategory::Subagent => "subagent",
ToolCategory::Introspection => "introspection",
ToolCategory::Config => "config",
}.to_string()
}
+163
View File
@@ -0,0 +1,163 @@
//! `ToolDiscovery` — records every tool actually offered to the LLM so the
//! approval / Security-groups UI can list and gate tools injected dynamically
//! outside the [`ToolRegistry`](crate::tools::ToolRegistry).
//!
//! Many tools reach the LLM outside the registry: `InterfaceTool` closures
//! (`write_todos`, `activate_tools`, `notify`, `show_file_to_user`, …), plugin
//! tools (Telegram's `send_voice_message`), and provider tools (`image_generate`,
//! memory backends). They are all assembled in one place —
//! [`AgentRunConfig::all_tool_defs`](crate::session::handler) — which is
//! the single point this service taps. Observing there *cannot drift* from what
//! is really offered, and covers every source uniformly, with zero per-component
//! wiring and no tool-name knowledge in core/plugins.
//!
//! See `docs/approval` and `docs/tools.md`.
use std::collections::HashSet;
use std::sync::{Arc, RwLock};
use serde_json::Value;
use sqlx::SqlitePool;
use tracing::warn;
use crate::db::known_tools;
pub struct ToolDiscovery {
db: Arc<SqlitePool>,
/// Names already persisted in this process. Keeps the per-round `observe`
/// call a cheap no-op after each tool's first sighting; DB writes happen
/// only for genuinely-new names.
seen: Arc<RwLock<HashSet<String>>>,
}
/// A tool extracted from an OpenAI function definition, ready to persist.
struct Draft {
name: String,
description: String,
schema: Option<String>,
}
impl ToolDiscovery {
pub fn new(db: Arc<SqlitePool>) -> Self {
Self { db, seen: Arc::new(RwLock::new(HashSet::new())) }
}
/// Observe the full OpenAI tool array for one round. Cheap no-op once every
/// name is known; otherwise persists the new tools in a background task so
/// the turn is never blocked on a DB write.
pub fn observe(&self, defs: &[Value]) {
let new_tools: Vec<Draft> = {
let seen = self.seen.read().unwrap();
defs.iter().filter_map(|d| Self::extract(d, &seen)).collect()
};
if new_tools.is_empty() {
return;
}
let db = Arc::clone(&self.db);
let seen = Arc::clone(&self.seen);
tokio::spawn(async move {
for t in new_tools {
match known_tools::upsert(&db, &t.name, &t.description, t.schema.as_deref()).await {
// Mark seen only after a successful write, so a transient DB
// error is retried on a later round instead of lost until restart.
Ok(()) => { seen.write().unwrap().insert(t.name); }
Err(e) => warn!(tool = %t.name, error = %e, "tool_discovery: failed to persist known tool"),
}
}
});
}
/// Pull `(name, description, parameters)` out of an OpenAI function
/// definition, skipping unnamed tools and ones already recorded.
fn extract(def: &Value, seen: &HashSet<String>) -> Option<Draft> {
let f = def.get("function")?;
let name = f.get("name")?.as_str()?.to_string();
if name.is_empty() || seen.contains(&name) {
return None;
}
let description = f.get("description").and_then(Value::as_str).unwrap_or("").to_string();
let schema = f.get("parameters").map(Value::to_string);
Some(Draft { name, description, schema })
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn def(name: &str) -> Value {
json!({
"type": "function",
"function": {
"name": name,
"description": format!("desc {name}"),
"parameters": { "type": "object" }
}
})
}
#[test]
fn extract_skips_unnamed_and_already_seen() {
let mut seen = HashSet::new();
seen.insert("write_todos".to_string());
// Already recorded → skipped.
assert!(ToolDiscovery::extract(&def("write_todos"), &seen).is_none());
// Missing function.name → skipped.
assert!(ToolDiscovery::extract(&json!({"type":"function","function":{}}), &seen).is_none());
// New, named → extracted with name/description/schema.
let d = ToolDiscovery::extract(&def("send_voice_message"), &seen).unwrap();
assert_eq!(d.name, "send_voice_message");
assert_eq!(d.description, "desc send_voice_message");
assert!(d.schema.is_some());
}
fn temp_db_path(tag: &str) -> String {
let mut p = std::env::temp_dir();
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
p.push(format!("skald-test-{tag}-{}-{nanos}.db", std::process::id()));
p.to_string_lossy().into_owned()
}
fn cleanup(path: &str) {
for suffix in ["", "-wal", "-shm"] {
let _ = std::fs::remove_file(format!("{path}{suffix}"));
}
}
/// The `observe` DB write happens in a spawned task; poll until it lands.
async fn wait_for_rows(pool: &SqlitePool, n: usize) -> Vec<known_tools::KnownTool> {
for _ in 0..100 {
let rows = known_tools::all(pool).await.unwrap();
if rows.len() >= n {
return rows;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
panic!("timed out waiting for {n} known_tools rows");
}
#[tokio::test]
async fn observe_persists_new_tools_and_dedups() {
let path = temp_db_path("discovery");
let pool = Arc::new(crate::db::init_system_pool(&path).await.unwrap());
let disc = ToolDiscovery::new(Arc::clone(&pool));
disc.observe(&[def("show_file_to_user"), def("notify")]);
let rows = wait_for_rows(&pool, 2).await;
let names: Vec<&str> = rows.iter().map(|r| r.name.as_str()).collect();
assert!(names.contains(&"show_file_to_user"));
assert!(names.contains(&"notify"));
// Re-observing a seen tool plus a new one records only the new one.
disc.observe(&[def("notify"), def("send_attachment")]);
let rows = wait_for_rows(&pool, 3).await;
assert_eq!(rows.len(), 3, "already-seen tools must not create duplicate rows");
pool.close().await;
cleanup(&path);
}
}
@@ -0,0 +1,156 @@
use std::collections::HashSet;
use std::sync::{Arc, RwLock};
use anyhow::Result;
use serde_json::{Value, json};
use sqlx::SqlitePool;
use crate::mcp::McpManager;
use crate::tools::tool_names::CONFIG_GROUP;
use crate::tools::{Tool, ToolDescriptionLength, truncate_label, MAX_LABEL_SHORT};
/// Per-session (or per-stack) tool that activates **tool groups** on demand.
///
/// A group is either:
/// - an **MCP server name** — loads that server's tools, or
/// - the reserved keyword `"config"` — loads all built-in `Config`-category
/// tools (system configuration: MCP/plugin/cron management, secrets).
///
/// When the LLM calls `activate_tools(["gmail", "config"])`:
/// - The in-memory grant set is updated immediately, so the group's tools appear
/// in the *next LLM round* of the current turn (via `all_tool_defs()`).
/// - If `stack_id` is `None` (root agent): grants are persisted to
/// `session_mcp_grants` — they survive across turns and restarts.
/// - If `stack_id` is `Some(id)` (sub-agent): grants are persisted to
/// `stack_mcp_grants` for that stack frame — they survive restarts but are
/// deleted when the frame terminates (`dispatch_call_agent` calls
/// `stack_mcp_grants::delete_for_stack` on cleanup).
///
/// The `session_mcp_grants` / `stack_mcp_grants` tables store the group string
/// verbatim, so `"config"` is persisted just like an MCP server name.
///
/// Not in the global `ToolRegistry` — injected as an `InterfaceTool` in
/// `build_agent_config` (root) and `dispatch_call_agent` (sub-agents).
pub struct ActivateTools {
pub pool: Arc<SqlitePool>,
pub session_id: i64,
/// `None` for root agents (session-scoped grants).
/// `Some(stack_id)` for sub-agents (stack-scoped grants, deleted on frame exit).
pub stack_id: Option<i64>,
pub mcp: Arc<McpManager>,
/// Shared in-memory grant set. Updated in-place on every call so subsequent
/// rounds within the same turn see the new tools via `all_tool_defs()`.
pub active_mcp_grants: Arc<RwLock<HashSet<String>>>,
}
impl Tool for ActivateTools {
fn name(&self) -> &str { crate::tools::tool_names::ACTIVATE_TOOLS }
fn category(&self) -> crate::tools::ToolCategory { crate::tools::ToolCategory::Config }
fn description(&self) -> &str {
"Activate one or more tool groups so their tools become available. \
A group is either an MCP server name (see the MCP list) or the reserved \
keyword `config`, which loads all system-configuration tools (managing \
MCP servers, plugins, scheduled cron jobs, and secrets). \
Pass an array of group names. \
Once activated, the tools are available from the next tool-call round onward."
}
fn parameters_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"groups": {
"type": "array",
"items": { "type": "string" },
"description": "Tool groups to activate: MCP server names and/or the reserved \
keyword \"config\" (e.g. [\"gmail\", \"config\"])."
}
},
"required": ["groups"]
})
}
fn describe(&self, args: &Value, _length: ToolDescriptionLength) -> String {
let names = args["groups"]
.as_array()
.map(|a| a.iter().filter_map(|v| v.as_str()).collect::<Vec<_>>().join(", "))
.unwrap_or_else(|| "?".to_string());
truncate_label(&format!("activate tools [{names}]"), MAX_LABEL_SHORT)
}
fn execute(&self, args: Value) -> Result<String> {
let names: Vec<String> = args["groups"]
.as_array()
.ok_or_else(|| anyhow::anyhow!("activate_tools: `groups` must be an array"))?
.iter()
.filter_map(|v| v.as_str().map(|s| s.to_string()))
.collect();
if names.is_empty() {
anyhow::bail!("activate_tools: `groups` is empty");
}
let available: HashSet<String> = self.mcp.tools()
.iter()
.map(|t| t.server_name.clone())
.collect();
let pool = Arc::clone(&self.pool);
let session_id = self.session_id;
let stack_id = self.stack_id;
let grants_set = Arc::clone(&self.active_mcp_grants);
// Persist to DB (session-scoped or stack-scoped) and update in-memory set.
// The reserved `config` group is stored verbatim, exactly like a server name.
tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(async {
for name in &names {
match stack_id {
None => {
crate::db::session_mcp_grants::grant(&pool, session_id, name).await?;
}
Some(sid) => {
crate::db::stack_mcp_grants::grant(&pool, sid, name).await?;
}
}
}
anyhow::Ok(())
})
})?;
// Update in-memory set so the next LLM round sees the new grants.
{
let mut set = grants_set.write()
.map_err(|_| anyhow::anyhow!("activate_tools: lock poisoned"))?;
for name in &names {
set.insert(name.clone());
}
}
let activated: Vec<String> = names.iter()
.map(|n| {
if n == CONFIG_GROUP {
// Built-in group — always available, no MCP server to reconnect.
format!("{n}")
} else if available.contains(n) {
format!("{n}")
} else {
format!("{n} (registered but not yet running — tools will appear after reconnect)")
}
})
.collect();
let scope = match stack_id {
None => "session".to_string(),
Some(s) => format!("stack {s}"),
};
Ok(format!(
"Tool groups activated for this {scope}: {}. \
Their tools are available from the next tool-call round.",
activated.join(", ")
))
}
}

Some files were not shown because too many files have changed in this diff Show More