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:
-98
@@ -1,98 +0,0 @@
|
||||
//! 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, which a stdout layer
|
||||
//! (see [`BootFormat`], wired in `main.rs`) renders cleanly — no timestamps,
|
||||
//! levels or targets, just the message, with failures shown in red.
|
||||
//!
|
||||
//! The same lines also 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.
|
||||
//!
|
||||
//! The stdout layer filters on the `boot` target only and is independent of
|
||||
//! `RUST_LOG`, so this output always appears regardless of the log filter.
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use tracing::field::{Field, Visit};
|
||||
use tracing::{Event, Level, info, warn};
|
||||
use tracing_subscriber::fmt::format::Writer;
|
||||
use tracing_subscriber::fmt::{FmtContext, FormatEvent, FormatFields};
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
|
||||
/// 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);
|
||||
}
|
||||
|
||||
/// Minimal stdout formatter for bootstrap lines: prints just the event's
|
||||
/// message (no timestamp/level/target), in red when the level is WARN or ERROR.
|
||||
pub struct BootFormat;
|
||||
|
||||
#[derive(Default)]
|
||||
struct MessageVisitor(String);
|
||||
|
||||
impl Visit for MessageVisitor {
|
||||
fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
|
||||
if field.name() == "message" {
|
||||
use std::fmt::Write;
|
||||
let _ = write!(self.0, "{value:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, N> FormatEvent<S, N> for BootFormat
|
||||
where
|
||||
S: tracing::Subscriber + for<'a> LookupSpan<'a>,
|
||||
N: for<'a> FormatFields<'a> + 'static,
|
||||
{
|
||||
fn format_event(
|
||||
&self,
|
||||
_ctx: &FmtContext<'_, S, N>,
|
||||
mut writer: Writer<'_>,
|
||||
event: &Event<'_>,
|
||||
) -> fmt::Result {
|
||||
let mut visitor = MessageVisitor::default();
|
||||
event.record(&mut visitor);
|
||||
|
||||
// Level ordering in tracing: TRACE > DEBUG > INFO > WARN > ERROR, so
|
||||
// `<= WARN` matches both WARN and ERROR.
|
||||
let is_failure = *event.metadata().level() <= Level::WARN;
|
||||
if writer.has_ansi_escapes() && is_failure {
|
||||
write!(writer, "\u{1b}[31m{}\u{1b}[0m", visitor.0)?;
|
||||
} else {
|
||||
write!(writer, "{}", visitor.0)?;
|
||||
}
|
||||
writeln!(writer)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
//! How this binary renders the bootstrap lines that `skald_core::boot` emits.
|
||||
//!
|
||||
//! Rendering is the shell's business, not the core's: a desktop bundle or a
|
||||
//! setup wizard would format the same `boot` target differently, or not at all.
|
||||
//! Wired in `main.rs` as a stdout layer filtered on `boot::TARGET`, independent
|
||||
//! of `RUST_LOG`, so this output always appears whatever the log filter is.
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use tracing::field::{Field, Visit};
|
||||
use tracing::{Event, Level};
|
||||
use tracing_subscriber::fmt::format::Writer;
|
||||
use tracing_subscriber::fmt::{FmtContext, FormatEvent, FormatFields};
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
|
||||
/// Minimal stdout formatter for bootstrap lines: prints just the event's
|
||||
/// message (no timestamp/level/target), in red when the level is WARN or ERROR.
|
||||
pub struct BootFormat;
|
||||
|
||||
#[derive(Default)]
|
||||
struct MessageVisitor(String);
|
||||
|
||||
impl Visit for MessageVisitor {
|
||||
fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
|
||||
if field.name() == "message" {
|
||||
use std::fmt::Write;
|
||||
let _ = write!(self.0, "{value:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, N> FormatEvent<S, N> for BootFormat
|
||||
where
|
||||
S: tracing::Subscriber + for<'a> LookupSpan<'a>,
|
||||
N: for<'a> FormatFields<'a> + 'static,
|
||||
{
|
||||
fn format_event(
|
||||
&self,
|
||||
_ctx: &FmtContext<'_, S, N>,
|
||||
mut writer: Writer<'_>,
|
||||
event: &Event<'_>,
|
||||
) -> fmt::Result {
|
||||
let mut visitor = MessageVisitor::default();
|
||||
event.record(&mut visitor);
|
||||
|
||||
// Level ordering in tracing: TRACE > DEBUG > INFO > WARN > ERROR, so
|
||||
// `<= WARN` matches both WARN and ERROR.
|
||||
let is_failure = *event.metadata().level() <= Level::WARN;
|
||||
if writer.has_ansi_escapes() && is_failure {
|
||||
write!(writer, "\u{1b}[31m{}\u{1b}[0m", visitor.0)?;
|
||||
} else {
|
||||
write!(writer, "{}", visitor.0)?;
|
||||
}
|
||||
writeln!(writer)
|
||||
}
|
||||
}
|
||||
+4
-4
@@ -4,7 +4,7 @@ use anyhow::{Context, Result};
|
||||
use serde::Deserialize;
|
||||
|
||||
pub use core_api::provider::LlmStrength;
|
||||
pub use crate::core::config::{
|
||||
pub use skald_core::config::{
|
||||
LlmConfig, TicConfig, CronConfig,
|
||||
CompactionConfig, DatetimeConfig, LlmRequestsLogConfig,
|
||||
};
|
||||
@@ -48,10 +48,10 @@ pub struct WebConfig {
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn into_split(self) -> (crate::core::config::CoreConfig, crate::frontend::config::FrontendConfig) {
|
||||
pub fn into_split(self) -> (skald_core::config::CoreConfig, crate::frontend::config::FrontendConfig) {
|
||||
let tz = self.timezone.clone();
|
||||
(
|
||||
crate::core::config::CoreConfig {
|
||||
skald_core::config::CoreConfig {
|
||||
llm: self.llm,
|
||||
tic: self.tic,
|
||||
cron: self.cron,
|
||||
@@ -74,7 +74,7 @@ impl Config {
|
||||
if !config_path.exists() {
|
||||
std::fs::copy(default_path, config_path)
|
||||
.with_context(|| format!("Failed to copy {DEFAULT_CONFIG} to {CONFIG}"))?;
|
||||
crate::boot::section(format!("Created {CONFIG} from {DEFAULT_CONFIG}"));
|
||||
skald_core::boot::section(format!("Created {CONFIG} from {DEFAULT_CONFIG}"));
|
||||
}
|
||||
|
||||
let content = std::fs::read_to_string(config_path)
|
||||
|
||||
@@ -1,259 +0,0 @@
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::{debug, trace, warn};
|
||||
|
||||
use crate::config::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
@@ -1,3 +0,0 @@
|
||||
pub use core_api::bus::{
|
||||
BusEvent, ChatEvent, ChatEventBus, ChatEventRole, CompactionEvent, RecvError, ToolCallEvent,
|
||||
};
|
||||
@@ -1,167 +0,0 @@
|
||||
//! 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");
|
||||
}
|
||||
}
|
||||
@@ -1,843 +0,0 @@
|
||||
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::core::approval::ApprovalManager;
|
||||
use crate::core::cron::TaskManager;
|
||||
use crate::core::db::{chat_history, chat_llm_tools, chat_sessions, chat_sessions_stack, config, sources};
|
||||
use crate::core::events::{GlobalEvent, ServerEvent};
|
||||
use crate::core::notification::Notification;
|
||||
use crate::core::session::handler::{ChatSessionHandler, InterfaceTool, PendingMsg, PendingUserInput};
|
||||
use crate::core::session::manager::ChatSessionManager;
|
||||
use crate::core::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::core::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::core::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::core::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::core::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(¬es).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
|
||||
}
|
||||
}
|
||||
@@ -1,162 +0,0 @@
|
||||
//! 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::core::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
|
||||
}
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
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,
|
||||
};
|
||||
@@ -1,107 +0,0 @@
|
||||
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::core::events::{GlobalEvent, ServerEvent};
|
||||
use crate::core::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;
|
||||
}
|
||||
}
|
||||
@@ -1,115 +0,0 @@
|
||||
//! 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 })
|
||||
}
|
||||
}
|
||||
@@ -1,512 +0,0 @@
|
||||
//! 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::core::chat_event_bus::{ChatEventBus, CompactionEvent};
|
||||
use crate::core::chatbot::ChatOptions;
|
||||
use crate::core::config::CompactionConfig;
|
||||
use crate::core::db::{chat_history, chat_llm_tools, chat_summaries};
|
||||
use crate::core::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::core::chatbot::LlmTurn::Message(resp) => resp.content,
|
||||
crate::core::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..])
|
||||
}
|
||||
@@ -1,118 +0,0 @@
|
||||
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>,
|
||||
}
|
||||
@@ -1,43 +0,0 @@
|
||||
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(())
|
||||
}
|
||||
}
|
||||
@@ -1,695 +0,0 @@
|
||||
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::core::chat_hub::ChatHub;
|
||||
use crate::core::db::chat_sessions;
|
||||
use crate::core::db::scheduled_jobs::{self, ScheduledJob};
|
||||
use crate::core::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::core::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::core::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::core::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::core::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::core::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::core::db::chat_history::append(
|
||||
pool, stack.id, &crate::core::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::core::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::core::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::core::session::handler::InterfaceTool {
|
||||
use crate::core::session::handler::{InterfaceTool, ToolFuture};
|
||||
use serde_json::json;
|
||||
|
||||
InterfaceTool {
|
||||
definition: json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": crate::core::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::core::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(())
|
||||
}
|
||||
@@ -1,97 +0,0 @@
|
||||
use anyhow::Result;
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
use crate::core::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(())
|
||||
}
|
||||
@@ -1,285 +0,0 @@
|
||||
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()
|
||||
}
|
||||
|
||||
pub async fn set_model_db_id(pool: &SqlitePool, id: i64, model_db_id: i64) -> anyhow::Result<()> {
|
||||
sqlx::query("UPDATE chat_history SET model_db_id = ? WHERE id = ?")
|
||||
.bind(model_db_id)
|
||||
.bind(id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
@@ -1,144 +0,0 @@
|
||||
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 }
|
||||
}
|
||||
@@ -1,76 +0,0 @@
|
||||
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,
|
||||
}))
|
||||
}
|
||||
@@ -1,190 +0,0 @@
|
||||
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::core::db::init_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);
|
||||
}
|
||||
}
|
||||
@@ -1,64 +0,0 @@
|
||||
//! 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,
|
||||
}))
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
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(())
|
||||
}
|
||||
@@ -1,103 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,94 +0,0 @@
|
||||
//! `known_tools` — every tool ever offered to the LLM, captured at injection
|
||||
//! time by [`crate::core::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::core::db::init_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);
|
||||
}
|
||||
}
|
||||
@@ -1,73 +0,0 @@
|
||||
//! 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::core::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)) => {}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -1,125 +0,0 @@
|
||||
//! DB operations for the `llm_requests` table.
|
||||
//!
|
||||
//! Every `chat_with_tools` call is logged here by the
|
||||
//! [`crate::core::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)
|
||||
}
|
||||
@@ -1,114 +0,0 @@
|
||||
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 }
|
||||
}
|
||||
@@ -1,129 +0,0 @@
|
||||
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(())
|
||||
}
|
||||
@@ -1,567 +0,0 @@
|
||||
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 anyhow::{Context, Result};
|
||||
use sqlx::{SqlitePool, sqlite::{SqliteConnectOptions, SqliteJournalMode, SqliteSynchronous}};
|
||||
use std::str::FromStr;
|
||||
use std::time::Duration;
|
||||
|
||||
/// System database: instance-wide state, shared by every user.
|
||||
///
|
||||
/// Fixed path — not configurable. Sibling `database/{userid}.db` files hold
|
||||
/// per-user content; keeping them in one directory makes backup, export and
|
||||
/// per-user erasure a matter of files rather than tables.
|
||||
pub const SYSTEM_DB_PATH: &str = "database/system.db";
|
||||
|
||||
pub async fn init_pool(path: &str) -> Result<SqlitePool> {
|
||||
// `create_if_missing` creates the file, never its parent directory.
|
||||
if let Some(parent) = std::path::Path::new(path).parent()
|
||||
&& !parent.as_os_str().is_empty()
|
||||
{
|
||||
std::fs::create_dir_all(parent)
|
||||
.with_context(|| format!("failed to create database directory {}", parent.display()))?;
|
||||
}
|
||||
let opts = SqliteConnectOptions::from_str(path)?
|
||||
.create_if_missing(true)
|
||||
// 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).
|
||||
.journal_mode(SqliteJournalMode::Wal)
|
||||
.synchronous(SqliteSynchronous::Normal)
|
||||
.busy_timeout(Duration::from_secs(5));
|
||||
let pool = SqlitePool::connect_with(opts).await?;
|
||||
create_tables(&pool).await?;
|
||||
crate::boot::section("Database initialised".to_string());
|
||||
Ok(pool)
|
||||
}
|
||||
|
||||
async fn create_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?;
|
||||
|
||||
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,
|
||||
model_db_id INTEGER REFERENCES llm_models(id),
|
||||
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 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 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 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 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 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 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 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 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 config (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL,
|
||||
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 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 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 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?;
|
||||
|
||||
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 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 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?;
|
||||
|
||||
// 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). See docs/approval + docs/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?;
|
||||
|
||||
// User directory + auth material. Read before every login, so it lives in
|
||||
// the system DB — 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(())
|
||||
}
|
||||
@@ -1,49 +0,0 @@
|
||||
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())
|
||||
}
|
||||
@@ -1,149 +0,0 @@
|
||||
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(())
|
||||
}
|
||||
@@ -1,103 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,211 +0,0 @@
|
||||
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(())
|
||||
}
|
||||
@@ -1,26 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,36 +0,0 @@
|
||||
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())
|
||||
}
|
||||
@@ -1,47 +0,0 @@
|
||||
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))
|
||||
}
|
||||
@@ -1,36 +0,0 @@
|
||||
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(())
|
||||
}
|
||||
@@ -1,85 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,540 +0,0 @@
|
||||
//! `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_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_pool_creates_the_database_directory() {
|
||||
let path = temp_db_path("users-mkdir");
|
||||
assert!(!std::path::Path::new(&path).parent().unwrap().exists());
|
||||
|
||||
let pool = crate::core::db::init_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::core::db::init_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::core::db::init_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::core::db::init_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::core::db::init_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>"));
|
||||
}
|
||||
}
|
||||
@@ -1,266 +0,0 @@
|
||||
//! 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::core::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::core::events::{GlobalEvent, ServerEvent};
|
||||
use crate::core::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);
|
||||
}
|
||||
}
|
||||
@@ -1,3 +0,0 @@
|
||||
pub use core_api::events::{
|
||||
ClientMessage, GlobalEvent, InboundDataMessage, ServerEvent,
|
||||
};
|
||||
@@ -1,103 +0,0 @@
|
||||
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,
|
||||
}
|
||||
}
|
||||
@@ -1,296 +0,0 @@
|
||||
/// 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::core::llm::LlmProviderRecord;
|
||||
use crate::core::llm::db as llm_db;
|
||||
use crate::core::provider::ProviderRegistry;
|
||||
use crate::core::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::core::tools::image_generate::ImageGenerateProvidersList { mgr: Arc::clone(&self) }) as Arc<dyn Tool>,
|
||||
Arc::new(crate::core::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;
|
||||
}
|
||||
}
|
||||
@@ -1,37 +0,0 @@
|
||||
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>,
|
||||
}
|
||||
@@ -1,94 +0,0 @@
|
||||
/// 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)
|
||||
}
|
||||
}
|
||||
@@ -1,142 +0,0 @@
|
||||
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::core::approval::{ApprovalManager, PendingApprovalInfo};
|
||||
use crate::core::clarification::{ClarificationManager, PendingClarificationInfo};
|
||||
use crate::core::elicitation::{ElicitationManager, ElicitationOutcome, PendingElicitationInfo};
|
||||
use crate::core::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
|
||||
}
|
||||
}
|
||||
@@ -1,681 +0,0 @@
|
||||
//! `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");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn composite_hash_is_deterministic_for_same_contents() {
|
||||
// Use the test file itself as a stand-in dependency: it exists on disk
|
||||
// and its contents are stable for the duration of the test.
|
||||
let me = Path::new(file!());
|
||||
let deps_a = vec![me.to_path_buf()];
|
||||
let deps_b = vec![me.to_path_buf()];
|
||||
assert_eq!(
|
||||
composite_hash_of(&deps_a).await.unwrap(),
|
||||
composite_hash_of(&deps_b).await.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn composite_hash_is_order_independent() {
|
||||
let me = Path::new(file!());
|
||||
let cargo = Path::new("Cargo.toml");
|
||||
let a = vec![me.to_path_buf(), cargo.to_path_buf()];
|
||||
let b = vec![cargo.to_path_buf(), me.to_path_buf()];
|
||||
assert_eq!(
|
||||
composite_hash_of(&a).await.unwrap(),
|
||||
composite_hash_of(&b).await.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[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);
|
||||
}
|
||||
}
|
||||
@@ -1,10 +0,0 @@
|
||||
//! 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};
|
||||
@@ -1,314 +0,0 @@
|
||||
use anyhow::{Context, Result};
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
use crate::config::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 (llm_requests.model_db_id references it),
|
||||
// only soft-deleted via `removed_at`. 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,
|
||||
}
|
||||
}
|
||||
@@ -1,582 +0,0 @@
|
||||
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::core::chatbot::ChatbotClient;
|
||||
use crate::core::chatbot::logging::{LoggingChatbotClient, LogSaveFlags};
|
||||
use crate::config::LlmStrength;
|
||||
use crate::core::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
|
||||
}
|
||||
}
|
||||
@@ -1,82 +0,0 @@
|
||||
pub(crate) mod db;
|
||||
pub mod manager;
|
||||
pub mod providers;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::core::chatbot::ChatbotClient;
|
||||
use crate::core::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,
|
||||
}
|
||||
@@ -1,119 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
|
||||
use crate::core::chatbot::anthropic::AnthropicClient;
|
||||
use crate::core::llm::{LlmModelRecord, LlmProviderRecord};
|
||||
use crate::core::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
|
||||
use crate::core::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 },
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,146 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
|
||||
use crate::core::chatbot::openai::OpenAiClient;
|
||||
use crate::core::llm::{LlmModelRecord, LlmProviderRecord};
|
||||
use crate::core::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
|
||||
use crate::core::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 },
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,86 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
|
||||
use crate::core::chatbot::lm_studio::LmStudioClient;
|
||||
use crate::core::llm::{LlmModelRecord, LlmProviderRecord};
|
||||
use crate::core::llm::providers::RemoteLlmModelInfo;
|
||||
use crate::core::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 },
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,39 +0,0 @@
|
||||
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::core::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),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,123 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
|
||||
use crate::core::chatbot::ollama::OllamaClient;
|
||||
use crate::core::llm::{LlmModelRecord, LlmProviderRecord};
|
||||
use crate::core::llm::providers::RemoteLlmModelInfo;
|
||||
use crate::core::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 },
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,99 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
|
||||
use crate::core::chatbot::openai::OpenAiClient;
|
||||
use crate::core::llm::{LlmModelRecord, LlmProviderRecord};
|
||||
use crate::core::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
|
||||
use crate::core::transcribe::TranscribeModelRecord;
|
||||
use crate::core::transcribe::openai_audio::OpenAiAudioTranscriber;
|
||||
use crate::core::tts::TtsModelRecord;
|
||||
use crate::core::tts::openai_tts::OpenAiTtsSynthesiser;
|
||||
use crate::core::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::core::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::core::tts::TextToSpeech>)
|
||||
})())
|
||||
}
|
||||
|
||||
fn build_transcriber(&self, record: &LlmProviderRecord, model: &TranscribeModelRecord) -> Option<Result<Arc<dyn crate::core::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::core::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 },
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,215 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
|
||||
use crate::core::chatbot::openai::OpenAiClient;
|
||||
use crate::core::image_generate::ImageGenerateModelRecord;
|
||||
use crate::core::image_generate::openrouter_image::OpenRouterImageGenerator;
|
||||
use crate::core::llm::{LlmModelRecord, LlmProviderRecord};
|
||||
use crate::core::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
|
||||
use crate::core::transcribe::TranscribeModelRecord;
|
||||
use crate::core::transcribe::openai_audio::OpenAiAudioTranscriber;
|
||||
use crate::core::tts::TtsModelRecord;
|
||||
use crate::core::tts::openai_tts::OpenAiTtsSynthesiser;
|
||||
use crate::core::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::core::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::core::tts::TextToSpeech>)
|
||||
})())
|
||||
}
|
||||
|
||||
fn build_transcriber(&self, record: &LlmProviderRecord, model: &TranscribeModelRecord) -> Option<Result<Arc<dyn crate::core::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::core::transcribe::Transcribe>)
|
||||
})())
|
||||
}
|
||||
|
||||
fn build_image_generator(&self, record: &LlmProviderRecord, model: &ImageGenerateModelRecord) -> Option<Result<Arc<dyn crate::core::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::core::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 },
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,145 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
|
||||
use crate::core::chatbot::openai::OpenAiClient;
|
||||
use crate::core::llm::{LlmModelRecord, LlmProviderRecord};
|
||||
use crate::core::llm::providers::{RemoteLlmModelInfo, extra_with_reasoning};
|
||||
use crate::core::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 +0,0 @@
|
||||
pub use core_api::location::{GpsCoord, LocationEntry, LocationManager, LocationUpdater};
|
||||
@@ -1,189 +0,0 @@
|
||||
//! 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;
|
||||
}
|
||||
}
|
||||
@@ -1,457 +0,0 @@
|
||||
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::core::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::core::db::mcp_events::insert(&pool, &source, &method, ¶ms).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::core::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::core::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::core::db::mcp_servers::UpsertParams<'_>) -> Result<Vec<String>> {
|
||||
let name = p.name.to_string();
|
||||
|
||||
crate::core::db::mcp_servers::upsert(&self.pool, p).await?;
|
||||
|
||||
let rows = crate::core::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::core::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::core::db::mcp_servers::set_enabled(&self.pool, name, enabled).await
|
||||
}
|
||||
|
||||
pub async fn list(&self) -> Result<Vec<McpServerInfo>> {
|
||||
let rows = crate::core::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!(" ({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",
|
||||
}
|
||||
}
|
||||
@@ -1,101 +0,0 @@
|
||||
//! 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::core::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() }
|
||||
}
|
||||
@@ -1,37 +0,0 @@
|
||||
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 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;
|
||||
@@ -1,28 +0,0 @@
|
||||
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,
|
||||
}
|
||||
@@ -1,140 +0,0 @@
|
||||
//! 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
|
||||
}
|
||||
}
|
||||
@@ -1,344 +0,0 @@
|
||||
pub use core_api::plugin::{Plugin, PluginContext, RouterFactory};
|
||||
pub use plugin_comfyui::ComfyUIPlugin;
|
||||
#[cfg(feature = "whisper-local")]
|
||||
pub use plugin_transcribe_whisper_local::WhisperLocalPlugin;
|
||||
pub use plugin_telegram_bot::TelegramPlugin;
|
||||
|
||||
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::core::db::plugins as db;
|
||||
use crate::core::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::core::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)
|
||||
}
|
||||
|
||||
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}"))
|
||||
}
|
||||
}
|
||||
@@ -1,108 +0,0 @@
|
||||
pub mod tickets;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::Result;
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
use crate::core::db::projects::{self, Project};
|
||||
use crate::core::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
|
||||
}
|
||||
@@ -1,178 +0,0 @@
|
||||
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::core::cron::TaskManager;
|
||||
use crate::core::db::{project_tickets, project_tickets::ProjectTicket, projects};
|
||||
use crate::core::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(())
|
||||
}
|
||||
}
|
||||
@@ -1,77 +0,0 @@
|
||||
// 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() });
|
||||
}
|
||||
}
|
||||
@@ -1,376 +0,0 @@
|
||||
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::core::db::tool_permission_groups::ToolPermissionGroup;
|
||||
use crate::core::approval::{ApprovalManager, RuleAction};
|
||||
use crate::core::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::core::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::core::db::tool_permission_groups::list(&self.db).await
|
||||
}
|
||||
|
||||
pub async fn get_group(&self, id: &str) -> Result<Option<ToolPermissionGroup>> {
|
||||
crate::core::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::core::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::core::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::core::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::core::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();
|
||||
}
|
||||
}
|
||||
@@ -1,65 +0,0 @@
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -1,9 +0,0 @@
|
||||
use crate::core::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;
|
||||
}
|
||||
@@ -1,329 +0,0 @@
|
||||
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::core::db::{chat_history, chat_llm_tools, chat_sessions_stack, scratchpad, stack_mcp_grants};
|
||||
use crate::core::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::core::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::core::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::core::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::core::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))
|
||||
}
|
||||
}
|
||||
@@ -1,115 +0,0 @@
|
||||
use serde_json::Value;
|
||||
use tracing::debug;
|
||||
|
||||
use super::ChatSessionHandler;
|
||||
use super::emitter::TurnEmitter;
|
||||
use crate::core::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::core::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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,205 +0,0 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::core::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::core::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::core::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::core::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::core::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::core::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::core::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,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,149 +0,0 @@
|
||||
//! 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::core::events::ServerEvent;
|
||||
use crate::core::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"})));
|
||||
}
|
||||
}
|
||||
@@ -1,155 +0,0 @@
|
||||
//! 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::core::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;
|
||||
}
|
||||
}
|
||||
@@ -1,138 +0,0 @@
|
||||
//! 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::core::approval::GateResult;
|
||||
use crate::core::db::chat_llm_tools;
|
||||
use crate::core::run_context::RunContext;
|
||||
use crate::core::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(¬e);
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,168 +0,0 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::core::mcp::McpManager;
|
||||
use crate::core::tools::Tool;
|
||||
use crate::core::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::core::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::core::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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,138 +0,0 @@
|
||||
//! 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::core::chatbot::{ChatOptions, LlmTurn};
|
||||
use crate::core::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()
|
||||
}
|
||||
@@ -1,433 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, info, trace};
|
||||
|
||||
use crate::core::tools::tool_names as tn;
|
||||
use crate::core::chat_event_bus::ToolCallEvent;
|
||||
use crate::core::chatbot::{LlmTurn, ToolCall};
|
||||
use crate::core::db::{chat_history, chat_llm_tools};
|
||||
use crate::core::events::ServerEvent;
|
||||
use crate::core::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::core::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?;
|
||||
chat_history::set_model_db_id(pool, message_id, cur_llm.model_db_id).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?;
|
||||
chat_history::set_model_db_id(pool, message_id, cur_llm.model_db_id).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::core::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)
|
||||
}
|
||||
}
|
||||
@@ -1,567 +0,0 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::{Value, json};
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
use crate::core::compactor::{ContextCompactor, SUMMARY_PREFIX};
|
||||
use crate::core::config::DatetimeConfig;
|
||||
use crate::core::db::{chat_history, chat_llm_tools, chat_summaries};
|
||||
use crate::core::mcp::McpManager;
|
||||
use crate::core::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::core::agents::load_prompt(agent_id)?;
|
||||
|
||||
let meta = crate::core::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::core::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::core::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)")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -1,677 +0,0 @@
|
||||
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::core::approval::ApprovalManager;
|
||||
use crate::core::run_context::RunContext;
|
||||
use crate::core::tools::tool_names as tn;
|
||||
use crate::core::chat_event_bus::{ChatEvent, ChatEventBus, ChatEventRole};
|
||||
use crate::core::clarification::ClarificationManager;
|
||||
use crate::core::compactor::ContextCompactor;
|
||||
use crate::core::config::DatetimeConfig;
|
||||
use crate::core::db::{chat_history, chat_sessions_stack};
|
||||
use crate::core::events::ServerEvent;
|
||||
use core_api::message_meta::MessageMetadata;
|
||||
use crate::core::llm::LlmManager;
|
||||
use crate::core::mcp::McpManager;
|
||||
use crate::core::image_generate::ImageGeneratorManager;
|
||||
use crate::core::memory::MemoryManager;
|
||||
use crate::core::tool_discovery::ToolDiscovery;
|
||||
use crate::core::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::core::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))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,92 +0,0 @@
|
||||
//! 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::core::chat_event_bus::ToolCallEvent;
|
||||
use crate::core::db::chat_llm_tools;
|
||||
use crate::core::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::core::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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,401 +0,0 @@
|
||||
use serde_json::Value;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::core::db::{chat_history, chat_llm_tools, chat_sessions_stack};
|
||||
use crate::core::events::ServerEvent;
|
||||
use crate::core::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::core::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 ¤t_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::core::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));
|
||||
}
|
||||
}
|
||||
@@ -1,180 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::core::approval::ApprovalManager;
|
||||
use crate::core::chat_event_bus::ChatEventBus;
|
||||
use crate::core::clarification::ClarificationManager;
|
||||
use crate::core::compactor::ContextCompactor;
|
||||
use crate::core::config::DatetimeConfig;
|
||||
use crate::core::db::{chat_sessions, chat_sessions_stack};
|
||||
use crate::core::llm::LlmManager;
|
||||
use crate::core::mcp::McpManager;
|
||||
use crate::core::image_generate::ImageGeneratorManager;
|
||||
use crate::core::memory::MemoryManager;
|
||||
use crate::core::run_context::{RunContext, RunContextManager};
|
||||
use crate::core::tool_discovery::ToolDiscovery;
|
||||
use crate::core::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)
|
||||
}
|
||||
}
|
||||
@@ -1,2 +0,0 @@
|
||||
pub mod handler;
|
||||
pub mod manager;
|
||||
@@ -1,96 +0,0 @@
|
||||
//! 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. Promote this block to a `SkaldApi` trait if/
|
||||
//! when `src/core/` is lifted into its own crate.
|
||||
|
||||
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::core::approval::ApprovalManager;
|
||||
use crate::core::chat_event_bus::ChatEventBus;
|
||||
use crate::core::chat_hub::ChatHub;
|
||||
use crate::core::clarification::ClarificationManager;
|
||||
use crate::core::command::LlmCommandManager;
|
||||
use crate::core::config_store::GlobalConfigManager;
|
||||
use crate::core::cron::TaskManager;
|
||||
use crate::core::elicitation::ElicitationManager;
|
||||
use crate::core::image_generate::ImageGeneratorManager;
|
||||
use crate::core::inbox::Inbox;
|
||||
use crate::core::latex::LatexCompiler;
|
||||
use crate::core::llm::LlmManager;
|
||||
use crate::core::location::LocationManager;
|
||||
use crate::core::mcp::McpManager;
|
||||
use crate::core::memory::MemoryManager;
|
||||
use crate::core::plugin::PluginManager;
|
||||
use crate::core::projects::tickets::ProjectTicketManager;
|
||||
use crate::core::projects::ProjectManager;
|
||||
use crate::core::provider::ProviderRegistry;
|
||||
use crate::core::run_context::RunContextManager;
|
||||
use crate::core::secrets::SecretsStore;
|
||||
use crate::core::session::manager::ChatSessionManager;
|
||||
use crate::core::tic::TicManager;
|
||||
use crate::core::tool_catalog::ToolCatalog;
|
||||
use crate::core::tools::ToolRegistry;
|
||||
use crate::core::transcribe::TranscribeManager;
|
||||
use crate::core::tts::TtsManager;
|
||||
|
||||
use super::Skald;
|
||||
|
||||
impl Skald {
|
||||
// Runtime / cross-cutting
|
||||
pub(crate) fn db(&self) -> &Arc<SqlitePool> { &self.rt.db }
|
||||
pub fn config(&self) -> &Arc<GlobalConfigManager> { &self.rt.config }
|
||||
pub fn config_properties(&self) -> &[core_api::ConfigSet] { &self.rt.config_properties }
|
||||
pub(crate) fn system_bus(&self) -> &Arc<SystemEventBus> { &self.rt.system_bus }
|
||||
pub(crate) 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 }
|
||||
}
|
||||
@@ -1,404 +0,0 @@
|
||||
//! 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::core::approval::ApprovalManager;
|
||||
use crate::core::chat_hub::ChatHub;
|
||||
use crate::core::clarification::ClarificationManager;
|
||||
use crate::core::command::LlmCommandManager;
|
||||
use crate::core::compactor::ContextCompactor;
|
||||
use crate::core::config::{CoreConfig, DatetimeConfig};
|
||||
use crate::core::cron::TaskManager;
|
||||
use crate::core::elicitation::ElicitationManager;
|
||||
use crate::core::image_generate::ImageGeneratorManager;
|
||||
use crate::core::inbox::Inbox;
|
||||
use crate::core::latex::LatexCompiler;
|
||||
use crate::core::llm::LlmManager;
|
||||
use crate::core::location::LocationManager;
|
||||
use crate::core::mcp::McpManager;
|
||||
use crate::core::memory::MemoryManager;
|
||||
use crate::core::plugin::PluginManager;
|
||||
use crate::core::projects::tickets::ProjectTicketManager;
|
||||
use crate::core::projects::ProjectManager;
|
||||
use crate::core::provider::ProviderRegistry;
|
||||
use crate::core::run_context::RunContextManager;
|
||||
use crate::core::secrets::SecretsStore;
|
||||
use crate::core::session::handler::{DEFAULT_MAX_PARALLEL_SUBAGENTS, DEFAULT_MAX_TOOL_ROUNDS};
|
||||
use crate::core::session::manager::ChatSessionManager;
|
||||
use crate::core::tic::TicManager;
|
||||
use crate::core::tool_catalog::ToolCatalog;
|
||||
use crate::core::tool_discovery::ToolDiscovery;
|
||||
use crate::core::tools::ToolRegistry;
|
||||
use crate::core::transcribe::TranscribeManager;
|
||||
use crate::core::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::core::llm::providers::openai::OpenAiProvider);
|
||||
provider_registry.register_builtin(crate::core::llm::providers::anthropic::AnthropicProvider::new());
|
||||
provider_registry.register_builtin(crate::core::llm::providers::openrouter::OpenRouterProvider::new());
|
||||
provider_registry.register_builtin(crate::core::llm::providers::ollama::OllamaProvider::new());
|
||||
provider_registry.register_builtin(crate::core::llm::providers::lm_studio::LmStudioProvider::new());
|
||||
provider_registry.register_builtin(crate::core::llm::providers::deepseek::DeepSeekProvider::new());
|
||||
provider_registry.register_builtin(crate::core::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::core::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::core::tools::fs::register_all(&mut tool_registry);
|
||||
tool_registry.register(crate::core::tools::ast_outline::AstOutline::new());
|
||||
tool_registry.register(crate::core::tools::exec::ExecuteCmd);
|
||||
tool_registry.register(crate::core::tools::read_notification::ReadNotification);
|
||||
tool_registry.register(crate::core::tools::restart::Restart);
|
||||
// Unified listing / toggling across mcp, plugins, cron (+ agents for list).
|
||||
tool_registry.register(crate::core::tools::list_items::ListItems::new(
|
||||
Arc::clone(&integrations.mcp), Arc::clone(&integrations.plugin_manager), Arc::clone(&tasks.cron)));
|
||||
tool_registry.register(crate::core::tools::toggle_item::ToggleItem::new(
|
||||
Arc::clone(&integrations.mcp), Arc::clone(&integrations.plugin_manager), Arc::clone(&tasks.cron)));
|
||||
tool_registry.register(crate::core::tools::register_mcp::RegisterMcp::new(Arc::clone(&integrations.mcp)));
|
||||
tool_registry.register(crate::core::tools::register_mcp::DeleteMcp::new(Arc::clone(&integrations.mcp)));
|
||||
tool_registry.register(crate::core::tools::cron_jobs::DeleteCronJob(Arc::clone(&tasks.cron)));
|
||||
tool_registry.register(crate::core::tools::set_secret::SetSecret(Arc::clone(&models.secrets)));
|
||||
tool_registry.register(crate::core::tools::list_secrets::ListSecrets(Arc::clone(&models.secrets)));
|
||||
tool_registry.register(crate::core::tools::configure_plugin::ConfigurePlugin(Arc::clone(&integrations.plugin_manager)));
|
||||
|
||||
// Mobile-connector control tools (plugin.md §11). The plugin Arc itself
|
||||
// implements `RelayAgent`; we look it up and bind the tools to it. The tools
|
||||
// call into the plugin lazily, so registering them before the plugin's
|
||||
// runloop starts is fine (they fail gracefully while stopped).
|
||||
if let Some(mc) = integrations.plugin_manager
|
||||
.get_plugin_typed::<plugin_mobile_connector::MobileConnectorPlugin>("mobile-connector")
|
||||
{
|
||||
let agent: Arc<dyn plugin_mobile_connector::RelayAgent> = mc;
|
||||
for tool in plugin_mobile_connector::mobile_tools(agent) {
|
||||
tool_registry.register_arc(tool);
|
||||
}
|
||||
info!("mobile-connector 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)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,96 +0,0 @@
|
||||
//! `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;
|
||||
}
|
||||
}
|
||||
@@ -1,63 +0,0 @@
|
||||
//! 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::core::chat_event_bus::ChatEventBus;
|
||||
use crate::core::config_store::GlobalConfigManager;
|
||||
|
||||
use super::supervisor::TaskSupervisor;
|
||||
|
||||
pub(super) struct Runtime {
|
||||
pub(super) db: Arc<SqlitePool>,
|
||||
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 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,
|
||||
config,
|
||||
config_properties: vec![crate::core::tic::config_set()],
|
||||
system_bus,
|
||||
event_bus,
|
||||
global_tx,
|
||||
shutdown_token: CancellationToken::new(),
|
||||
supervisor: TaskSupervisor::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,59 +0,0 @@
|
||||
//! 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");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,101 +0,0 @@
|
||||
//! 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::core::config::CoreConfig;
|
||||
use crate::core::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::core::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");
|
||||
}
|
||||
@@ -1,253 +0,0 @@
|
||||
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::core::chat_hub::ChatHub;
|
||||
use crate::core::config::TicConfig;
|
||||
use crate::core::config_store::GlobalConfigManager;
|
||||
use crate::core::db::mcp_events;
|
||||
use crate::core::run_context::{RunContext, RunContextManager};
|
||||
use crate::core::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::core::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::core::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}")
|
||||
}
|
||||
@@ -1,119 +0,0 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::core::mcp::McpManager;
|
||||
use crate::core::tools::{ToolCategory, ToolDescriptionLength, ToolRegistry};
|
||||
use crate::core::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()
|
||||
}
|
||||
@@ -1,163 +0,0 @@
|
||||
//! `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::core::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::core::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::core::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::core::db::init_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);
|
||||
}
|
||||
}
|
||||
@@ -1,156 +0,0 @@
|
||||
use std::collections::HashSet;
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
use anyhow::Result;
|
||||
use serde_json::{Value, json};
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
use crate::core::mcp::McpManager;
|
||||
use crate::core::tools::tool_names::CONFIG_GROUP;
|
||||
use crate::core::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::core::tools::tool_names::ACTIVATE_TOOLS }
|
||||
|
||||
fn category(&self) -> crate::core::tools::ToolCategory { crate::core::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::core::db::session_mcp_grants::grant(&pool, session_id, name).await?;
|
||||
}
|
||||
Some(sid) => {
|
||||
crate::core::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
Reference in New Issue
Block a user