llm request tracking, user context cleanup, minor fixes
This commit is contained in:
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user