llm request tracking, user context cleanup, minor fixes

This commit is contained in:
2026-07-11 00:23:14 +01:00
parent 5936cf0b2e
commit 587958ffe0
24 changed files with 402 additions and 347 deletions
+29 -53
View File
@@ -1,11 +1,12 @@
//! 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).
//! [`LoggingChatbotClient`] intercepts every `chat_with_tools_raw` call, captures
//! the raw HTTP request/response from the inner provider, persists a **metadata-only**
//! row to `llm_requests` in `system.db` (fire-and-forget), then returns the raw data
//! to the caller so it can write the **payload** to the user's own database.
//!
//! The LLM loop is completely unaware of this: it only holds an
//! `Arc<dyn ChatbotClient>` and calls `chat_with_tools` as usual.
//! The split keeps conversation content (payloads) behind the user key while
//! metadata (cost, tokens, timing) stays in the admin-readable registry.
use std::sync::Arc;
use std::time::Instant;
@@ -21,26 +22,10 @@ use super::{ChatOptions, ChatResponse, ChatbotClient, LlmRawMeta, LlmTurn, Messa
// ─────────────────────────────────────────────────────────────────────────────
/// 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 {
@@ -48,9 +33,8 @@ impl LoggingChatbotClient {
inner: Arc<dyn ChatbotClient>,
pool: Arc<SqlitePool>,
model_name: impl Into<String>,
flags: LogSaveFlags,
) -> Self {
Self { inner, pool, model_name: model_name.into(), flags }
Self { inner, pool, model_name: model_name.into() }
}
}
@@ -65,20 +49,35 @@ impl ChatbotClient for LoggingChatbotClient {
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.
/// Passthrough that drops the raw meta. Used by callers that do not need
/// payload capture (e.g. the compactor).
async fn chat_with_tools(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<LlmTurn> {
let (turn, _) = self.chat_with_tools_raw(messages, tools, options).await?;
Ok(turn)
}
/// Intercepts the call, delegates to `inner.chat_with_tools_raw` to capture
/// HTTP wire data, writes a **metadata-only** row to `system.db`, then returns
/// the raw data so the caller can persist payloads to the user's own database.
async fn chat_with_tools_raw(
&self,
messages: &[Value],
tools: &[Value],
options: &ChatOptions,
) -> anyhow::Result<(LlmTurn, Option<LlmRawMeta>)> {
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 user_id = options.user_id.clone();
let request_id = options.request_id.clone();
let model_name = self.model_name.clone();
let pool = Arc::clone(&self.pool);
@@ -90,24 +89,13 @@ impl ChatbotClient for LoggingChatbotClient {
(*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 {
request_id,
user_id,
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),
@@ -119,7 +107,7 @@ impl ChatbotClient for LoggingChatbotClient {
}
});
Ok(turn)
Ok((turn, meta))
}
Err(e) => {
@@ -127,13 +115,11 @@ impl ChatbotClient for LoggingChatbotClient {
tokio::spawn(async move {
if let Err(log_err) = llm_requests::insert(&pool, llm_requests::LlmRequestRow {
request_id,
user_id,
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,
@@ -149,14 +135,4 @@ impl ChatbotClient for LoggingChatbotClient {
}
}
}
/// 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
}
}
+2
View File
@@ -313,6 +313,8 @@ impl ContextCompactor {
temperature: Some(0.3),
session_id: Some(session_id),
stack_id: Some(stack_id),
user_id: None,
request_id: None,
};
let turn = llm.client.chat_with_tools(&messages_payload, &[], &options).await
@@ -0,0 +1,32 @@
//! DB operations for the `llm_request_payloads` table (owner bucket).
//!
//! Full request/response payloads + headers for each LLM call. Lives in
//! `{userid}.db` (encrypted), correlated with the metadata row in `system.db`
//! via `request_id`.
use anyhow::Result;
use sqlx::SqlitePool;
pub struct PayloadRow {
pub request_json: String,
pub request_headers: Option<String>,
pub response_json: Option<String>,
pub response_headers: Option<String>,
pub request_id: String,
}
pub async fn insert(pool: &SqlitePool, row: PayloadRow) -> Result<()> {
sqlx::query(
"INSERT INTO llm_request_payloads
(request_id, request_json, request_headers, response_json, response_headers)
VALUES (?, ?, ?, ?, ?)",
)
.bind(&row.request_id)
.bind(&row.request_json)
.bind(&row.request_headers)
.bind(&row.response_json)
.bind(&row.response_headers)
.execute(pool)
.await?;
Ok(())
}
@@ -1,9 +1,9 @@
//! 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.
//! Deletes expired metadata rows according to the retention settings in
//! [`LlmRequestsLogConfig`], then `VACUUM`s to reclaim freed pages.
//! Payload/header nulling is gone — those columns moved to `llm_request_payloads`
//! in the owner bucket.
use std::sync::Arc;
use std::time::Duration;
@@ -31,27 +31,6 @@ pub fn spawn(
_ = 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"),
@@ -59,7 +38,7 @@ pub fn spawn(
Err(e) => warn!(error = %e, "llm_requests: delete old rows failed"),
}
}
// VACUUM reclaims pages freed by DELETE/UPDATE NULL.
// VACUUM reclaims pages freed by DELETE.
match sqlx::query("VACUUM").execute(&*pool).await {
Ok(_) => info!("llm_requests: VACUUM complete"),
Err(e) => warn!(error = %e, "llm_requests: VACUUM failed"),
+9 -61
View File
@@ -1,7 +1,9 @@
//! DB operations for the `llm_requests` table.
//! DB operations for the `llm_requests` table (metadata only).
//!
//! Every `chat_with_tools` call is logged here by the
//! [`crate::chatbot::logging::LoggingChatbotClient`] wrapper.
//! Payloads (request/response bodies + headers) live in `llm_request_payloads`
//! in the owner bucket (`{userid}.db`), correlated by `request_id`.
//! Rows are retained for `llm.request_log.retention_days` days (default 14).
use anyhow::Result;
@@ -12,17 +14,11 @@ pub mod cleanup;
// ── Row struct ────────────────────────────────────────────────────────────────
pub struct LlmRequestRow {
pub request_id: Option<String>,
pub user_id: Option<String>,
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>,
@@ -40,21 +36,17 @@ pub struct LlmRequestRow {
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,
request_id, user_id, session_id, stack_id, model_name,
error_text, input_tokens, output_tokens, duration_ms,
cache_read_tokens, cache_creation_tokens
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
RETURNING id",
)
.bind(&row.request_id)
.bind(&row.user_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)
@@ -79,47 +71,3 @@ pub async fn delete_old_rows(pool: &SqlitePool, days: u32) -> Result<u64> {
.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)
}
+24 -4
View File
@@ -10,6 +10,7 @@ pub mod config;
pub mod job_runs;
pub mod known_tools;
pub mod llm_requests;
pub mod llm_request_payloads;
pub mod mcp_events;
pub mod mcp_servers;
pub mod plugins;
@@ -317,16 +318,17 @@ async fn create_registry_tables(pool: &SqlitePool) -> Result<()> {
// decrypting anything: the admin sees how much, when and which model — never
// what was said. `session_id` / `stack_id` are bare integers, not foreign
// keys, precisely because the rows they point at live in another file.
// `user_id` correlates the row with the payload in `{userid}.db`.
// Payloads (request/response bodies, headers) live in `llm_request_payloads`
// in the owner bucket — they are conversation content, behind the user key.
sqlx::query(
"CREATE TABLE IF NOT EXISTS llm_requests (
id INTEGER PRIMARY KEY AUTOINCREMENT,
request_id TEXT,
user_id TEXT,
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,
@@ -691,6 +693,23 @@ pub async fn create_owner_tables(pool: &SqlitePool) -> Result<()> {
.execute(pool)
.await?;
// Full request/response payloads for telemetry. Lives in the owner bucket
// (per-user, encrypted) because it is conversation content. Correlated with
// the metadata row in `system.db` via `request_id` (uuid).
sqlx::query(
"CREATE TABLE IF NOT EXISTS llm_request_payloads (
id INTEGER PRIMARY KEY AUTOINCREMENT,
request_id TEXT NOT NULL,
request_json TEXT NOT NULL DEFAULT '',
request_headers TEXT,
response_json TEXT,
response_headers TEXT,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
)",
)
.execute(pool)
.await?;
Ok(())
}
@@ -744,6 +763,7 @@ mod tests {
one("INSERT INTO secrets (key, value) VALUES ('k', 'v')").await.unwrap();
one("INSERT INTO projects (id, name, path) VALUES (1, 'p', '/tmp')").await.unwrap();
one("INSERT INTO project_tickets (project_id, title, job_id) VALUES (1, 't', 1)").await.unwrap();
one("INSERT INTO llm_request_payloads (request_id, request_json) VALUES ('r1', '{}')").await.unwrap();
pool.close().await;
let _ = std::fs::remove_dir_all(&dir);
+2
View File
@@ -456,6 +456,8 @@ mod tests {
async fn cleartext_user_round_trips_with_and_without_a_verifier() {
let path = temp_db_path("users-clear");
let pool = crate::db::init_system_pool(&path).await.unwrap();
crate::db::roles::insert(&pool, "children", "Children", "default", None)
.await.unwrap();
insert(&pool, "u-1", "kid", None, "children", &cleartext()).await.unwrap();
insert(&pool, "u-2", "kiosk", None, "children", &Credentials::Cleartext(None)).await.unwrap();
+13 -13
View File
@@ -9,7 +9,7 @@ use tokio::sync::RwLock;
use tracing::{info, warn};
use crate::chatbot::ChatbotClient;
use crate::chatbot::logging::{LoggingChatbotClient, LogSaveFlags};
use crate::chatbot::logging::LoggingChatbotClient;
use core_api::provider::LlmStrength;
use crate::provider::{ApiProvider, ProviderRegistry, ReasoningMode};
@@ -69,15 +69,15 @@ pub struct LlmManager {
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>,
/// When `true`, every LLM entry is wrapped with [`LoggingChatbotClient`].
log_enabled: bool,
}
impl LlmManager {
pub async fn new(
pool: Arc<SqlitePool>,
registry: Arc<ProviderRegistry>,
log_flags: Option<LogSaveFlags>,
pool: Arc<SqlitePool>,
registry: Arc<ProviderRegistry>,
log_enabled: bool,
) -> Result<Arc<Self>> {
let mgr = Arc::new(Self {
pool,
@@ -89,7 +89,7 @@ impl LlmManager {
}),
catalog: RwLock::new(HashMap::new()),
model_meta_cache: RwLock::new(HashMap::new()),
log_flags,
log_enabled,
});
mgr.reload().await?;
Ok(mgr)
@@ -461,9 +461,9 @@ impl LlmManager {
}
};
let log_config = self.log_flags.map(|f| (Arc::clone(&self.pool), f));
let log_pool = self.log_enabled.then(|| Arc::clone(&self.pool));
let entry = match build_entry(&self.registry, &provider, &model, model.id, log_config) {
let entry = match build_entry(&self.registry, &provider, &model, model.id, log_pool) {
Ok(e) => Arc::new(e),
Err(e) => {
warn!(model = %model.name, error = %e, "failed to build LLM entry, skipping");
@@ -505,7 +505,7 @@ fn build_entry(
provider: &LlmProviderRecord,
model: &LlmModelRecord,
model_db_id: i64,
log_config: Option<(Arc<SqlitePool>, LogSaveFlags)>,
log_pool: Option<Arc<SqlitePool>>,
) -> Result<LlmEntry> {
let built = registry.get(&provider.provider)
.ok_or_else(|| anyhow::anyhow!("unknown provider type '{}'", provider.provider))?
@@ -516,9 +516,9 @@ fn build_entry(
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,
let client: Arc<dyn ChatbotClient> = match log_pool {
Some(pool) => Arc::new(LoggingChatbotClient::new(inner, pool, &model.name)),
None => inner,
};
Ok(LlmEntry {
-3
View File
@@ -96,10 +96,7 @@ impl PluginManager {
.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 _,
@@ -13,6 +13,7 @@ use tokio_util::sync::CancellationToken;
use tracing::{error, warn};
use crate::chatbot::{ChatOptions, LlmTurn};
use crate::db::llm_request_payloads;
use crate::llm::{LlmEntry, LlmStrength};
use super::ChatSessionHandler;
@@ -54,12 +55,15 @@ impl ChatSessionHandler {
let mut tried_this_round: Vec<String> = vec![cur_name.clone()];
loop {
let request_id = uuid::Uuid::new_v4().to_string();
let options = ChatOptions {
model: cur_llm.model.clone(),
max_tokens: None,
temperature: None,
session_id: Some(self.session_id),
stack_id: Some(stack_id),
user_id: Some(self.user_id.clone()),
request_id: Some(request_id.clone()),
};
// Clone the Arc so the in-flight future does not borrow `cur_llm` across
@@ -68,13 +72,33 @@ impl ChatSessionHandler {
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,
r = client.chat_with_tools_raw(messages.as_slice(), tool_defs, &options) => r,
};
let e = match call_result {
Ok(t) => {
Ok((turn, meta)) => {
self.llm_manager.mark_success(cur_name).await;
return RoundLlm::Turn(t);
// Persist the payload (request/response bodies + headers) to the
// user's own database. Fire-and-forget — a failed write must not
// break the turn. The metadata row is already written by the
// logging wrapper to system.db with the same request_id.
if let Some(meta) = meta {
let pool = Arc::clone(&self.db);
let rid = request_id.clone();
tokio::spawn(async move {
let row = llm_request_payloads::PayloadRow {
request_id: rid,
request_json: meta.request_body.map(|v| v.to_string()).unwrap_or_default(),
request_headers: meta.request_headers.map(|v| v.to_string()),
response_json: meta.response_body.map(|v| v.to_string()),
response_headers: meta.response_headers.map(|v| v.to_string()),
};
if let Err(e) = llm_request_payloads::insert(&pool, row).await {
tracing::warn!(error = %e, "llm_request_payloads: failed to insert");
}
});
}
return RoundLlm::Turn(turn);
}
Err(e) => e,
};
@@ -262,6 +262,9 @@ impl ApprovalDecision {
pub struct ChatSessionHandler {
pub session_id: i64,
pub(super) db: Arc<SqlitePool>,
/// The authenticated user who owns this session. Threaded into `ChatOptions`
/// so the telemetry metadata row in `system.db` carries `user_id`.
pub(super) user_id: String,
pub(super) llm_manager: Arc<LlmManager>,
pub(super) max_history_messages: usize,
pub(super) max_tool_rounds: usize,
@@ -329,6 +332,7 @@ impl ChatSessionHandler {
pub fn new(
session_id: i64,
db: Arc<SqlitePool>,
user_id: String,
llm_manager: Arc<LlmManager>,
max_history_messages: usize,
max_tool_rounds: usize,
@@ -353,6 +357,7 @@ impl ChatSessionHandler {
Self {
session_id,
db,
user_id,
llm_manager,
max_history_messages,
max_tool_rounds,
+4
View File
@@ -22,6 +22,7 @@ use super::handler::ChatSessionHandler;
pub struct ChatSessionManager {
db: Arc<SqlitePool>,
user_id: String,
llm_manager: Arc<LlmManager>,
max_history_messages: usize,
max_tool_rounds: usize,
@@ -47,6 +48,7 @@ pub struct ChatSessionManager {
impl ChatSessionManager {
pub fn new(
db: Arc<SqlitePool>,
user_id: String,
llm_manager: Arc<LlmManager>,
max_history_messages: usize,
max_tool_rounds: usize,
@@ -66,6 +68,7 @@ impl ChatSessionManager {
) -> Self {
Self {
db,
user_id,
llm_manager,
max_history_messages,
max_tool_rounds,
@@ -152,6 +155,7 @@ impl ChatSessionManager {
let handler = Arc::new(ChatSessionHandler::new(
session_id,
self.db.clone(),
self.user_id.clone(),
Arc::clone(&self.llm_manager),
self.max_history_messages,
self.max_tool_rounds,
+3 -10
View File
@@ -72,16 +72,8 @@ impl Models {
let provider_registry = Arc::new(provider_registry);
info!("provider registry ready ({} built-in providers)", provider_registry.all().len());
let log_flags = config.llm.requests_log.as_ref().filter(|r| r.enabled).map(|r| {
use crate::chatbot::logging::LogSaveFlags;
LogSaveFlags {
request_payload: r.request_payload_save,
response_payload: r.response_payload_save,
request_headers: r.request_header_save,
response_headers: r.response_header_save,
}
});
let llm_manager = LlmManager::new(Arc::clone(&rt.db), Arc::clone(&provider_registry), log_flags).await?;
let log_enabled = config.llm.requests_log.as_ref().is_some_and(|r| r.enabled);
let llm_manager = LlmManager::new(Arc::clone(&rt.db), Arc::clone(&provider_registry), log_enabled).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");
@@ -346,6 +338,7 @@ impl Conversation {
let manager = Arc::new(ChatSessionManager::new(
Arc::clone(&rt.db),
String::new(),
Arc::clone(&models.llm_manager),
config.llm.max_history_messages,
config.llm.max_tool_rounds.unwrap_or(DEFAULT_MAX_TOOL_ROUNDS),
@@ -45,6 +45,7 @@ use crate::inbox::Inbox;
use crate::llm::LlmManager;
use crate::mcp::McpManager;
use crate::memory::MemoryManager;
use crate::projects::tickets::ProjectTicketManager;
use crate::run_context::RunContextManager;
use crate::session::handler::{DEFAULT_MAX_PARALLEL_SUBAGENTS, DEFAULT_MAX_TOOL_ROUNDS};
use crate::session::manager::ChatSessionManager;
@@ -62,6 +63,7 @@ pub struct UserContext {
pub sessions: Arc<ChatSessionManager>,
pub chat_hub: Arc<ChatHub>,
pub cron: Arc<TaskManager>,
pub tickets: Arc<ProjectTicketManager>,
pub approval: Arc<ApprovalManager>,
pub clarification: Arc<ClarificationManager>,
pub elicitation: Arc<ElicitationManager>,
@@ -154,6 +156,7 @@ impl UserContextFactory {
let manager = Arc::new(ChatSessionManager::new(
Arc::clone(&pool),
user_id.to_string(),
Arc::clone(&self.llm_manager),
self.max_history_messages,
self.max_tool_rounds,
@@ -189,12 +192,29 @@ impl UserContextFactory {
cron.set_self_arc(Arc::clone(&cron));
chat_hub.set_task_mgr(Arc::clone(&cron));
// Per-user ticket manager — wired to the per-user TaskManager so
// `start_ticket` spawns jobs in the user's own pool.
let tickets = ProjectTicketManager::new(Arc::clone(&pool));
tickets.set_task_manager(Arc::clone(&cron));
// Per-user cron loop. `start()` observes the shutdown token, so it stops on
// shutdown; adopting it lets the supervisor also join it. The name is leaked
// to satisfy the `&'static str` label — bounded by the (small) user count.
let name: &'static str = Box::leak(format!("cron:{user_id}").into_boxed_str());
self.supervisor.adopt(name, Arc::clone(&cron).start(self.shutdown_token.clone()));
// Per-user ticket-listener: reacts to JobCompleted events for this user's
// tickets. All users' listeners receive the event (global system bus); only
// the one that owns the ticket does the UPDATE — others no-op on 0 rows.
let tname: &'static str = Box::leak(format!("tickets:{user_id}").into_boxed_str());
self.supervisor.adopt_one(
tname,
Arc::clone(&tickets).start_listener(
Arc::clone(&self.system_bus),
self.shutdown_token.clone(),
),
);
Ok(Arc::new(UserContext {
user_id: user_id.to_string(),
pool,
@@ -202,6 +222,7 @@ impl UserContextFactory {
sessions: manager,
chat_hub,
cron,
tickets,
approval,
clarification,
elicitation,
+17 -38
View File
@@ -1,6 +1,11 @@
//! 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.
//!
//! Owner-bound background loops (cron, session-cancel, ticket-listener, tic) have
//! moved per-user into `UserContextFactory::build`. What remains here are the
//! instance-wide tasks: LLM-log cleanup on the registry pool, and MCP server
//! initialization.
use std::sync::Arc;
@@ -14,6 +19,10 @@ use super::runtime::Runtime;
/// Resolves the construction cycles (`cron ↔ session ↔ hub`, `ticket → cron`,
/// `mcp → elicitation`) via the managers' `OnceLock` setters.
///
/// These wire the **global** bundles, which are transitional and will be removed
/// once all call-sites are per-user (Phase 6). They remain constructed so any
/// not-yet-migrated accessor does not panic on a `None` OnceLock.
pub(super) fn wire(
tasks: &Tasks,
conversation: &Conversation,
@@ -29,14 +38,16 @@ pub(super) fn wire(
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.
/// Spawns the instance-wide background tasks.
///
/// Owner-bound loops (cron, session-cancel, ticket-listener, tic) are **not**
/// spawned here — they run per-user inside `UserContext`. Session cancellation is
/// handled directly by the API handlers (which have `AuthUser` and resolve the
/// per-user context). TIC is deferred until connectors return (§13).
pub(super) fn spawn_background(
rt: &Runtime,
tasks: &Tasks,
conversation: &Conversation,
_tasks: &Tasks,
_conversation: &Conversation,
integrations: &Integrations,
config: &CoreConfig,
) {
@@ -52,29 +63,6 @@ pub(super) fn spawn_background(
);
}
// 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
@@ -89,13 +77,4 @@ pub(super) fn spawn_background(
}
});
}
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");
}
+5
View File
@@ -507,6 +507,11 @@ mod tests {
let system = db::init_system_pool(dir.join("system.db").to_str().unwrap())
.await
.unwrap();
// Seed the non-admin role some tests use, so the FK on
// users.role_id → roles(id) holds.
db::roles::insert(&system, "children", "Children", "default", None)
.await
.unwrap();
let users = UserManager {
system: Arc::new(system),