hugo/RuDS
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

Implement persistent conversational AI chat

Browse Source
This commit is contained in:
Chili Palmer
2026-07-19 16:18:01 +02:00
parent fb5cae2131
commit ac611e3f7f
31 changed files with 4035 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/bds-core/migrations/20260718000000_initial_schema/up.sql
View File

@@ -199,9 +199,7 @@ CREATE TABLE IF NOT EXISTS chat_messages (
content TEXT,
tool_call_id TEXT,
tool_calls TEXT,
created_at INTEGER NOT NULL,
cache_read_tokens INTEGER,
cache_write_tokens INTEGER
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS ai_providers (

3
crates/bds-core/migrations/20260719141036_add_chat_cache_token_usage/down.sql Normal file
View File

@@ -0,0 +1,3 @@
-- This file should undo anything in `up.sql`
ALTER TABLE chat_messages DROP COLUMN cache_write_tokens;
ALTER TABLE chat_messages DROP COLUMN cache_read_tokens;

3
crates/bds-core/migrations/20260719141036_add_chat_cache_token_usage/up.sql Normal file
View File

@@ -0,0 +1,3 @@
-- Your SQL goes here
ALTER TABLE chat_messages ADD COLUMN cache_read_tokens INTEGER;
ALTER TABLE chat_messages ADD COLUMN cache_write_tokens INTEGER;

50
crates/bds-core/src/db/migrations.rs
View File

@@ -36,7 +36,7 @@ mod tests {
let applied = db
.conn()
.with_migrations(|conn| conn.applied_migrations().unwrap().len());
assert_eq!(applied, 4);
assert_eq!(applied, 5);
}
#[test]
@@ -156,4 +156,52 @@ mod tests {
assert_eq!(model_ref.as_deref(), Some("model-package"));
assert_eq!(usage_tokens, (Some(56), Some(78), Some(12), Some(34)));
}
#[test]
fn existing_chat_schema_is_upgraded_with_cache_token_columns() {
let db = Database::open_in_memory().unwrap();
db.conn().with_migrations(|conn| {
for _ in 0..4 {
conn.run_next_migration(MIGRATIONS).unwrap();
}
});
db.conn()
.with(|conn| {
diesel::insert_into(chat_conversations::table)
.values((
chat_conversations::id.eq("existing"),
chat_conversations::title.eq("Existing chat"),
chat_conversations::created_at.eq(1_i64),
chat_conversations::updated_at.eq(1_i64),
))
.execute(conn)?;
diesel::insert_into(chat_messages::table)
.values((
chat_messages::conversation_id.eq("existing"),
chat_messages::role.eq("assistant"),
chat_messages::created_at.eq(1_i64),
chat_messages::token_usage_input.eq(Some(8)),
chat_messages::token_usage_output.eq(Some(5)),
))
.execute(conn)
})
.unwrap();
run_migrations(db.conn()).unwrap();
let usage = db
.conn()
.with(|conn| {
chat_messages::table
.select((
chat_messages::token_usage_input,
chat_messages::token_usage_output,
chat_messages::cache_read_tokens,
chat_messages::cache_write_tokens,
))
.first::<(Option<i32>, Option<i32>, Option<i32>, Option<i32>)>(conn)
})
.unwrap();
assert_eq!(usage, (Some(8), Some(5), None, None));
}
}

134
crates/bds-core/src/db/queries/chat.rs Normal file
View File

@@ -0,0 +1,134 @@
use diesel::prelude::*;
use crate::db::DbConnection;
use crate::db::schema::{chat_conversations, chat_messages};
use crate::model::{ChatConversation, ChatMessage, NewChatConversation, NewChatMessage};
pub fn insert_conversation(
conn: &DbConnection,
conversation: &NewChatConversation<'_>,
) -> QueryResult<ChatConversation> {
conn.with(|connection| {
diesel::insert_into(chat_conversations::table)
.values(conversation)
.execute(connection)?;
chat_conversations::table
.find(conversation.id)
.select(ChatConversation::as_select())
.first(connection)
})
}
pub fn get_conversation(conn: &DbConnection, id: &str) -> QueryResult<ChatConversation> {
conn.with(|connection| {
chat_conversations::table
.find(id)
.select(ChatConversation::as_select())
.first(connection)
})
}
pub fn list_conversations(conn: &DbConnection) -> QueryResult<Vec<ChatConversation>> {
conn.with(|connection| {
chat_conversations::table
.order((
chat_conversations::updated_at.desc(),
chat_conversations::id.desc(),
))
.select(ChatConversation::as_select())
.load(connection)
})
}
pub fn rename_conversation(
conn: &DbConnection,
id: &str,
title: &str,
updated_at: i64,
) -> QueryResult<ChatConversation> {
conn.with(|connection| {
diesel::update(chat_conversations::table.find(id))
.set((
chat_conversations::title.eq(title),
chat_conversations::updated_at.eq(updated_at),
))
.execute(connection)?;
chat_conversations::table
.find(id)
.select(ChatConversation::as_select())
.first(connection)
})
}
pub fn set_conversation_model(
conn: &DbConnection,
id: &str,
model: &str,
updated_at: i64,
) -> QueryResult<usize> {
conn.with(|connection| {
diesel::update(chat_conversations::table.find(id))
.set((
chat_conversations::model.eq(model),
chat_conversations::updated_at.eq(updated_at),
))
.execute(connection)
})
}
pub fn set_session_id(
conn: &DbConnection,
id: &str,
session_id: Option<&str>,
updated_at: i64,
) -> QueryResult<usize> {
conn.with(|connection| {
diesel::update(chat_conversations::table.find(id))
.set((
chat_conversations::copilot_session_id.eq(session_id),
chat_conversations::updated_at.eq(updated_at),
))
.execute(connection)
})
}
pub fn delete_conversation(conn: &DbConnection, id: &str) -> QueryResult<usize> {
conn.with(|connection| {
connection.transaction(|connection| {
diesel::delete(chat_messages::table.filter(chat_messages::conversation_id.eq(id)))
.execute(connection)?;
diesel::delete(chat_conversations::table.find(id)).execute(connection)
})
})
}
pub fn insert_message(
conn: &DbConnection,
message: &NewChatMessage<'_>,
updated_at: i64,
) -> QueryResult<ChatMessage> {
conn.with(|connection| {
connection.transaction(|connection| {
diesel::insert_into(chat_messages::table)
.values(message)
.execute(connection)?;
diesel::update(chat_conversations::table.find(message.conversation_id))
.set(chat_conversations::updated_at.eq(updated_at))
.execute(connection)?;
chat_messages::table
.order(chat_messages::id.desc())
.select(ChatMessage::as_select())
.first(connection)
})
})
}
pub fn list_messages(conn: &DbConnection, conversation_id: &str) -> QueryResult<Vec<ChatMessage>> {
conn.with(|connection| {
chat_messages::table
.filter(chat_messages::conversation_id.eq(conversation_id))
.order((chat_messages::created_at.asc(), chat_messages::id.asc()))
.select(ChatMessage::as_select())
.load(connection)
})
}

1
crates/bds-core/src/db/queries/mod.rs
View File

@@ -1,3 +1,4 @@
pub mod chat;
pub mod db_notification;
pub mod generated_file_hash;
pub mod import_definition;

5
crates/bds-core/src/db/types.rs
View File

@@ -5,8 +5,8 @@ use diesel::sql_types::{Integer, Text};
use diesel::sqlite::{Sqlite, SqliteValue};
use crate::model::{
NotificationAction, NotificationEntity, PostStatus, ProposalKind, ProposalStatus, ScriptKind,
ScriptStatus, TemplateKind, TemplateStatus,
ChatRole, NotificationAction, NotificationEntity, PostStatus, ProposalKind, ProposalStatus,
ScriptKind, ScriptStatus, TemplateKind, TemplateStatus,
};
#[derive(Debug, AsExpression, FromSqlRow)]
@@ -95,3 +95,4 @@ text_enum_sql!(NotificationEntity);
text_enum_sql!(NotificationAction);
text_enum_sql!(ProposalKind);
text_enum_sql!(ProposalStatus);
text_enum_sql!(ChatRole);

977
crates/bds-core/src/engine/chat.rs Normal file
View File

@@ -0,0 +1,977 @@
use std::collections::{BTreeMap, HashMap};
use std::io::Read;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, OnceLock, mpsc};
use std::time::Duration;
use diesel::prelude::*;
use reqwest::blocking::Client;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use uuid::Uuid;
use crate::db::DbConnection as Connection;
use crate::db::queries::chat as queries;
use crate::db::schema::ai_models;
use crate::engine::ai::{self, AiEndpointConfig, TokenUsage};
use crate::engine::{EngineError, EngineResult, chat_tools};
use crate::model::{ChatConversation, ChatMessage, ChatRole, NewChatConversation, NewChatMessage};
use crate::util::now_unix_ms;
const DEFAULT_CONTEXT_TOKENS: usize = 32_768;
const DEFAULT_OUTPUT_TOKENS: u64 = 16_384;
const MAX_TOOL_ROUNDS: usize = 10;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatToolCall {
pub id: String,
pub name: String,
pub arguments: Value,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ChatModelInfo {
pub provider: String,
pub id: String,
pub name: String,
pub family: Option<String>,
pub context_window: u64,
pub max_output_tokens: u64,
pub supports_tools: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChatEvent {
Started {
conversation_id: String,
},
Content {
conversation_id: String,
content: String,
},
ToolStarted {
conversation_id: String,
name: String,
},
ToolFinished {
conversation_id: String,
name: String,
},
Finished {
conversation_id: String,
usage: TokenUsage,
},
Failed {
conversation_id: String,
message: String,
},
Cancelled {
conversation_id: String,
},
Navigate {
destination: String,
entity_id: Option<String>,
},
}
#[derive(Clone)]
pub struct ChatSendOptions {
pub endpoint: Option<AiEndpointConfig>,
pub model: Option<String>,
pub context_tokens: Option<usize>,
pub max_output_tokens: u64,
pub max_tool_rounds: usize,
pub enable_tools: bool,
pub event_handler: Option<Arc<dyn Fn(ChatEvent) + Send + Sync>>,
}
impl Default for ChatSendOptions {
fn default() -> Self {
Self {
endpoint: None,
model: None,
context_tokens: None,
max_output_tokens: DEFAULT_OUTPUT_TOKENS,
max_tool_rounds: MAX_TOOL_ROUNDS,
enable_tools: true,
event_handler: None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct ChatTurnResult {
pub content: String,
pub usage: TokenUsage,
pub cancelled: bool,
}
#[derive(Debug, Clone, PartialEq, Default)]
pub struct AssembledResponse {
pub content: String,
pub tool_calls: Vec<ChatToolCall>,
pub usage: TokenUsage,
pub session_id: Option<String>,
}
#[derive(Default)]
struct PartialToolCall {
id: String,
name: String,
arguments: String,
}
#[derive(Default)]
pub struct SseAssembler {
buffer: Vec<u8>,
content: String,
tools: BTreeMap<u64, PartialToolCall>,
usage: TokenUsage,
session_id: Option<String>,
done: bool,
}
impl SseAssembler {
pub fn feed(&mut self, bytes: &[u8]) -> EngineResult<()> {
self.buffer.extend_from_slice(bytes);
while let Some((end, separator_len)) = event_boundary(&self.buffer) {
let event = self.buffer.drain(..end).collect::<Vec<_>>();
self.buffer.drain(..separator_len);
self.process_event(&event)?;
}
Ok(())
}
pub fn snapshot(&self) -> &str {
&self.content
}
pub fn finish(mut self) -> EngineResult<AssembledResponse> {
if !self.buffer.is_empty() {
let trailing = std::mem::take(&mut self.buffer);
self.process_event(&trailing)?;
}
let tool_calls = self
.tools
.into_values()
.map(|tool| {
let arguments = if tool.arguments.trim().is_empty() {
json!({})
} else {
serde_json::from_str(&tool.arguments).map_err(|error| {
EngineError::Parse(format!(
"invalid arguments for tool {}: {error}",
tool.name
))
})?
};
Ok(ChatToolCall {
id: tool.id,
name: tool.name,
arguments,
})
})
.collect::<EngineResult<Vec<_>>>()?;
Ok(AssembledResponse {
content: self.content,
tool_calls,
usage: self.usage,
session_id: self.session_id,
})
}
fn process_event(&mut self, event: &[u8]) -> EngineResult<()> {
let text = std::str::from_utf8(event)
.map_err(|error| EngineError::Parse(format!("invalid SSE encoding: {error}")))?;
let data = text
.lines()
.filter_map(|line| line.strip_prefix("data:"))
.map(str::trim_start)
.collect::<Vec<_>>()
.join("\n");
if data.is_empty() || data == "[DONE]" {
self.done |= data == "[DONE]";
return Ok(());
}
let body: Value = serde_json::from_str(&data)
.map_err(|error| EngineError::Parse(format!("malformed SSE event: {error}")))?;
self.apply_json(&body)
}
fn apply_json(&mut self, body: &Value) -> EngineResult<()> {
if let Some(error) = body.get("error") {
return Err(EngineError::Parse(format!("provider error: {error}")));
}
if let Some(session_id) = body
.get("session_id")
.or_else(|| body.get("sessionId"))
.and_then(Value::as_str)
{
self.session_id = Some(session_id.to_string());
}
merge_usage(&mut self.usage, body);
for choice in body
.get("choices")
.and_then(Value::as_array)
.into_iter()
.flatten()
{
let delta = choice
.get("delta")
.or_else(|| choice.get("message"))
.unwrap_or(&Value::Null);
if let Some(content) = delta.get("content").and_then(Value::as_str) {
self.content.push_str(content);
}
for tool in delta
.get("tool_calls")
.and_then(Value::as_array)
.into_iter()
.flatten()
{
let index = tool.get("index").and_then(Value::as_u64).unwrap_or(0);
let partial = self.tools.entry(index).or_default();
if let Some(id) = tool.get("id").and_then(Value::as_str) {
partial.id.push_str(id);
}
let function = tool.get("function").unwrap_or(&Value::Null);
if let Some(name) = function.get("name").and_then(Value::as_str) {
partial.name.push_str(name);
}
if let Some(arguments) = function.get("arguments").and_then(Value::as_str) {
partial.arguments.push_str(arguments);
}
}
}
Ok(())
}
}
pub fn create_conversation(
conn: &Connection,
model: Option<&str>,
) -> EngineResult<ChatConversation> {
let model = model.filter(|value| !value.trim().is_empty());
let title = model
.map(|model| format!("Chat with {model}"))
.unwrap_or_else(|| "New Chat".to_string());
create_conversation_titled(conn, model, &title)
}
pub fn create_conversation_titled(
conn: &Connection,
model: Option<&str>,
title: &str,
) -> EngineResult<ChatConversation> {
let model = model.filter(|value| !value.trim().is_empty());
let title = title.trim();
if title.is_empty() {
return Err(EngineError::Validation(
"conversation title is required".to_string(),
));
}
let id = Uuid::new_v4().to_string();
let now = now_unix_ms();
Ok(queries::insert_conversation(
conn,
&NewChatConversation {
id: &id,
title,
model,
copilot_session_id: None,
created_at: now,
updated_at: now,
},
)?)
}
pub fn rename_conversation(
conn: &Connection,
id: &str,
title: &str,
) -> EngineResult<ChatConversation> {
let title = title.trim();
if title.is_empty() {
return Err(EngineError::Validation(
"conversation title is required".to_string(),
));
}
Ok(queries::rename_conversation(
conn,
id,
title,
now_unix_ms(),
)?)
}
pub fn set_conversation_model(conn: &Connection, id: &str, model: &str) -> EngineResult<()> {
let model = model.trim();
if model.is_empty() {
return Err(EngineError::Validation(
"chat model is required".to_string(),
));
}
if queries::set_conversation_model(conn, id, model, now_unix_ms())? == 0 {
return Err(EngineError::NotFound(format!("conversation {id}")));
}
Ok(())
}
pub fn list_conversations(conn: &Connection) -> EngineResult<Vec<ChatConversation>> {
Ok(queries::list_conversations(conn)?)
}
pub fn list_models(conn: &Connection) -> EngineResult<Vec<ChatModelInfo>> {
let rows = conn.with(|connection| {
ai_models::table
.select((
ai_models::provider,
ai_models::model_id,
ai_models::name,
ai_models::family,
ai_models::context_window,
ai_models::max_output_tokens,
ai_models::tool_call,
))
.order((ai_models::provider.asc(), ai_models::name.asc()))
.load::<(String, String, String, Option<String>, i32, i32, i32)>(connection)
})?;
Ok(rows
.into_iter()
.map(
|(provider, id, name, family, context_window, max_output_tokens, tool_call)| {
ChatModelInfo {
provider,
id,
name,
family,
context_window: context_window.max(0) as u64,
max_output_tokens: max_output_tokens.max(0) as u64,
supports_tools: tool_call != 0,
}
},
)
.collect())
}
pub fn get_conversation(conn: &Connection, id: &str) -> EngineResult<ChatConversation> {
queries::get_conversation(conn, id).map_err(|error| match error {
diesel::result::Error::NotFound => EngineError::NotFound(format!("conversation {id}")),
error => error.into(),
})
}
pub fn delete_conversation(conn: &Connection, id: &str) -> EngineResult<()> {
cancel_chat(id);
if queries::delete_conversation(conn, id)? == 0 {
return Err(EngineError::NotFound(format!("conversation {id}")));
}
Ok(())
}
pub fn list_messages(conn: &Connection, id: &str) -> EngineResult<Vec<ChatMessage>> {
Ok(queries::list_messages(conn, id)?)
}
pub fn insert_message(
conn: &Connection,
conversation_id: &str,
role: ChatRole,
content: Option<&str>,
tool_call_id: Option<&str>,
tool_calls: Option<&str>,
usage: TokenUsage,
) -> EngineResult<ChatMessage> {
let now = now_unix_ms();
Ok(queries::insert_message(
conn,
&NewChatMessage {
conversation_id,
role,
content,
tool_call_id,
tool_calls,
created_at: now,
cache_read_tokens: token_i32(usage.cache_read_tokens),
cache_write_tokens: token_i32(usage.cache_write_tokens),
token_usage_input: token_i32(usage.input_tokens),
token_usage_output: token_i32(usage.output_tokens),
},
now,
)?)
}
pub fn subscribe_events() -> mpsc::Receiver<ChatEvent> {
let (sender, receiver) = mpsc::channel();
listeners()
.lock()
.expect("chat listeners lock")
.push(sender);
receiver
}
pub fn cancel_chat(conversation_id: &str) -> bool {
let state = in_flight()
.lock()
.expect("chat cancellation lock")
.get(conversation_id)
.cloned();
if let Some(state) = state {
state.store(true, Ordering::SeqCst);
true
} else {
false
}
}
#[allow(clippy::too_many_arguments)]
pub fn send_chat_message(
conn: &Connection,
data_dir: &std::path::Path,
project_id: &str,
offline_mode: bool,
conversation_id: &str,
content: &str,
options: ChatSendOptions,
) -> EngineResult<ChatTurnResult> {
let content = content.trim();
if content.is_empty() {
return Err(EngineError::Validation(
"chat message is required".to_string(),
));
}
let conversation = get_conversation(conn, conversation_id)?;
let endpoint = match options.endpoint.clone() {
Some(endpoint) => validate_runtime_endpoint(endpoint)?,
None => ai::active_endpoint(conn, offline_mode)?,
};
let model = options
.model
.clone()
.or(conversation.model.clone())
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| endpoint.model.clone());
if model.trim().is_empty() {
return Err(EngineError::Validation(
"AI unavailable - configure a chat model in Settings".to_string(),
));
}
let cancelled = Arc::new(AtomicBool::new(false));
{
let mut active = in_flight().lock().expect("chat cancellation lock");
if active.contains_key(conversation_id) {
return Err(EngineError::Conflict(
"a response is already streaming for this conversation".to_string(),
));
}
active.insert(conversation_id.to_string(), Arc::clone(&cancelled));
}
let _guard = InFlightGuard(conversation_id.to_string());
emit(
&options,
ChatEvent::Started {
conversation_id: conversation_id.to_string(),
},
);
set_conversation_model(conn, conversation_id, &model)?;
insert_message(
conn,
conversation_id,
ChatRole::User,
Some(content),
None,
None,
TokenUsage::default(),
)?;
let result = run_turns(
conn,
data_dir,
project_id,
conversation_id,
&endpoint,
&model,
&options,
&cancelled,
);
match &result {
Ok(turn) if turn.cancelled => emit(
&options,
ChatEvent::Cancelled {
conversation_id: conversation_id.to_string(),
},
),
Ok(turn) => emit(
&options,
ChatEvent::Finished {
conversation_id: conversation_id.to_string(),
usage: turn.usage,
},
),
Err(error) => emit(
&options,
ChatEvent::Failed {
conversation_id: conversation_id.to_string(),
message: error.to_string(),
},
),
}
result
}
#[allow(clippy::too_many_arguments)]
fn run_turns(
conn: &Connection,
data_dir: &std::path::Path,
project_id: &str,
conversation_id: &str,
endpoint: &AiEndpointConfig,
model: &str,
options: &ChatSendOptions,
cancelled: &AtomicBool,
) -> EngineResult<ChatTurnResult> {
let mut total_usage = TokenUsage::default();
let max_rounds = options.max_tool_rounds.min(MAX_TOOL_ROUNDS);
let supports_tools = options.enable_tools && chat_tools::model_supports_tools(conn, model)?;
let catalog_model = list_models(conn)?
.into_iter()
.find(|candidate| candidate.id == model);
let context_window = options.context_tokens.unwrap_or_else(|| {
catalog_model
.as_ref()
.map(|model| model.context_window as usize)
.filter(|window| *window > 0)
.unwrap_or(DEFAULT_CONTEXT_TOKENS)
});
let max_output_tokens = catalog_model
.as_ref()
.map(|model| model.max_output_tokens)
.filter(|limit| *limit > 0)
.map_or(options.max_output_tokens, |limit| {
limit.min(options.max_output_tokens)
});
let tool_specs = supports_tools.then(chat_tools::tool_specs);
let tool_budget = tool_specs
.as_ref()
.map(|tools| approximate_tokens(&Value::Array(tools.clone())))
.unwrap_or(0);
let context_budget = context_window
.saturating_sub(max_output_tokens as usize)
.saturating_sub(tool_budget)
.max(1_024.min(context_window));
for round in 0..=max_rounds {
if cancelled.load(Ordering::SeqCst) {
return Ok(ChatTurnResult {
usage: total_usage,
cancelled: true,
..Default::default()
});
}
let messages = build_context(conn, project_id, conversation_id, context_budget)?;
let mut payload = json!({
"model": model,
"messages": messages,
"stream": true,
"stream_options": {"include_usage": true},
"max_tokens": max_output_tokens,
});
if let Some(tools) = tool_specs.as_ref() {
payload["tools"] = Value::Array(tools.clone());
payload["tool_choice"] = json!("auto");
}
let response = request_completion(endpoint, &payload, cancelled, |content| {
emit(
options,
ChatEvent::Content {
conversation_id: conversation_id.to_string(),
content: content.to_string(),
},
);
})?;
add_usage(&mut total_usage, response.usage);
if let Some(session_id) = response.session_id.as_deref() {
queries::set_session_id(conn, conversation_id, Some(session_id), now_unix_ms())?;
}
if cancelled.load(Ordering::SeqCst) {
if !response.content.is_empty() {
insert_message(
conn,
conversation_id,
ChatRole::Assistant,
Some(&response.content),
None,
None,
response.usage,
)?;
}
return Ok(ChatTurnResult {
content: response.content,
usage: total_usage,
cancelled: true,
});
}
if !response.tool_calls.is_empty() && round == max_rounds {
if !response.content.is_empty() {
insert_message(
conn,
conversation_id,
ChatRole::Assistant,
Some(&response.content),
None,
None,
response.usage,
)?;
}
return Err(EngineError::Validation(format!(
"chat exceeded the {max_rounds}-round tool limit"
)));
}
let serialized_calls = (!response.tool_calls.is_empty())
.then(|| serialize_tool_calls(&response.tool_calls))
.transpose()?;
insert_message(
conn,
conversation_id,
ChatRole::Assistant,
(!response.content.is_empty()).then_some(response.content.as_str()),
None,
serialized_calls.as_deref(),
response.usage,
)?;
if response.tool_calls.is_empty() {
return Ok(ChatTurnResult {
content: response.content,
usage: total_usage,
cancelled: false,
});
}
for (index, call) in response.tool_calls.iter().enumerate() {
if cancelled.load(Ordering::SeqCst) {
persist_cancelled_tool_results(
conn,
conversation_id,
&response.tool_calls[index..],
)?;
return Ok(ChatTurnResult {
usage: total_usage,
cancelled: true,
..Default::default()
});
}
emit(
options,
ChatEvent::ToolStarted {
conversation_id: conversation_id.to_string(),
name: call.name.clone(),
},
);
if cancelled.load(Ordering::SeqCst) {
persist_cancelled_tool_results(
conn,
conversation_id,
&response.tool_calls[index..],
)?;
return Ok(ChatTurnResult {
usage: total_usage,
cancelled: true,
..Default::default()
});
}
let result =
chat_tools::execute(conn, data_dir, project_id, &call.name, &call.arguments);
let result = match result {
Ok(result) => result,
Err(error) => json!({"success": false, "error": error.to_string()}),
};
if call.name == "navigate"
&& let Some(navigation) = result.get("navigation")
&& let Some(destination) = navigation.get("destination").and_then(Value::as_str)
{
emit(
options,
ChatEvent::Navigate {
destination: destination.to_string(),
entity_id: navigation
.get("entity_id")
.and_then(Value::as_str)
.map(str::to_string),
},
);
}
let result_text = serde_json::to_string(&result)?;
insert_message(
conn,
conversation_id,
ChatRole::Tool,
Some(&result_text),
Some(&call.id),
None,
TokenUsage::default(),
)?;
emit(
options,
ChatEvent::ToolFinished {
conversation_id: conversation_id.to_string(),
name: call.name.clone(),
},
);
}
}
unreachable!("bounded loop returns on its final iteration")
}
fn persist_cancelled_tool_results(
conn: &Connection,
conversation_id: &str,
calls: &[ChatToolCall],
) -> EngineResult<()> {
let result = json!({"success": false, "cancelled": true}).to_string();
for call in calls {
insert_message(
conn,
conversation_id,
ChatRole::Tool,
Some(&result),
Some(&call.id),
None,
TokenUsage::default(),
)?;
}
Ok(())
}
pub fn build_context(
conn: &Connection,
project_id: &str,
conversation_id: &str,
token_budget: usize,
) -> EngineResult<Vec<Value>> {
let system = chat_tools::system_prompt(conn, project_id)?;
let messages = list_messages(conn, conversation_id)?;
let mut groups: Vec<Vec<Value>> = Vec::new();
let mut current = Vec::new();
for message in messages {
let value = message_json(&message)?;
if message.role == ChatRole::System {
continue;
}
if message.role == ChatRole::User && !current.is_empty() {
groups.push(std::mem::take(&mut current));
}
current.push(value);
}
if !current.is_empty() {
groups.push(current);
}
let system_value = json!({"role": "system", "content": system});
let mut used = approximate_tokens(&system_value);
let mut selected = Vec::new();
for group in groups.into_iter().rev() {
let cost = group.iter().map(approximate_tokens).sum::<usize>();
if used + cost > token_budget && !selected.is_empty() {
continue;
}
used += cost;
selected.push(group);
if used >= token_budget {
break;
}
}
selected.reverse();
let mut result = vec![system_value];
result.extend(selected.into_iter().flatten());
Ok(result)
}
fn request_completion(
endpoint: &AiEndpointConfig,
payload: &Value,
cancelled: &AtomicBool,
mut on_content: impl FnMut(&str),
) -> EngineResult<AssembledResponse> {
let client = Client::builder()
.connect_timeout(Duration::from_secs(5))
.timeout(Duration::from_secs(300))
.build()?;
let mut request = client
.post(chat_completions_url(&endpoint.url))
.json(payload);
if let Some(api_key) = endpoint.api_key.as_deref() {
request = request.bearer_auth(api_key);
}
let mut response = request.send()?.error_for_status()?;
let is_stream = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| value.contains("text/event-stream"));
if !is_stream {
let body: Value = response.json()?;
let mut assembler = SseAssembler::default();
assembler.apply_json(&body)?;
return assembler.finish();
}
let mut assembler = SseAssembler::default();
let mut chunk = [0_u8; 4096];
loop {
if cancelled.load(Ordering::SeqCst) {
break;
}
let count = response.read(&mut chunk)?;
if count == 0 {
break;
}
let old_len = assembler.snapshot().len();
assembler.feed(&chunk[..count])?;
if assembler.snapshot().len() != old_len {
on_content(assembler.snapshot());
}
}
assembler.finish()
}
fn message_json(message: &ChatMessage) -> EngineResult<Value> {
let mut value = json!({"role": message.role.as_str()});
if let Some(content) = message.content.as_deref() {
value["content"] = json!(content);
}
if let Some(tool_call_id) = message.tool_call_id.as_deref() {
value["tool_call_id"] = json!(tool_call_id);
}
if let Some(tool_calls) = message.tool_calls.as_deref() {
value["tool_calls"] = serde_json::from_str(tool_calls)?;
}
Ok(value)
}
fn serialize_tool_calls(calls: &[ChatToolCall]) -> EngineResult<String> {
let calls = calls
.iter()
.map(|call| {
json!({
"id": call.id,
"type": "function",
"function": {
"name": call.name,
"arguments": call.arguments.to_string(),
}
})
})
.collect::<Vec<_>>();
Ok(serde_json::to_string(&calls)?)
}
fn approximate_tokens(value: &Value) -> usize {
value.to_string().chars().count().div_ceil(4).max(1)
}
fn merge_usage(usage: &mut TokenUsage, body: &Value) {
let source = body.get("usage").unwrap_or(&Value::Null);
usage.input_tokens = source
.get("prompt_tokens")
.and_then(Value::as_u64)
.or(usage.input_tokens);
usage.output_tokens = source
.get("completion_tokens")
.and_then(Value::as_u64)
.or(usage.output_tokens);
usage.cache_read_tokens = source
.get("prompt_tokens_details")
.and_then(|value| value.get("cached_tokens"))
.and_then(Value::as_u64)
.or_else(|| source.get("cache_read_tokens").and_then(Value::as_u64))
.or(usage.cache_read_tokens);
usage.cache_write_tokens = source
.get("completion_tokens_details")
.and_then(|value| value.get("cached_tokens"))
.and_then(Value::as_u64)
.or_else(|| source.get("cache_write_tokens").and_then(Value::as_u64))
.or(usage.cache_write_tokens);
}
fn add_usage(total: &mut TokenUsage, current: TokenUsage) {
total.input_tokens = add_optional(total.input_tokens, current.input_tokens);
total.output_tokens = add_optional(total.output_tokens, current.output_tokens);
total.cache_read_tokens = add_optional(total.cache_read_tokens, current.cache_read_tokens);
total.cache_write_tokens = add_optional(total.cache_write_tokens, current.cache_write_tokens);
}
fn add_optional(left: Option<u64>, right: Option<u64>) -> Option<u64> {
match (left, right) {
(None, None) => None,
(left, right) => Some(left.unwrap_or(0).saturating_add(right.unwrap_or(0))),
}
}
fn token_i32(value: Option<u64>) -> Option<i32> {
value.map(|value| i32::try_from(value).unwrap_or(i32::MAX))
}
fn validate_runtime_endpoint(endpoint: AiEndpointConfig) -> EngineResult<AiEndpointConfig> {
if endpoint.url.trim().is_empty() || endpoint.model.trim().is_empty() {
return Err(EngineError::Validation(
"AI unavailable - configure endpoint and model in Settings".to_string(),
));
}
Ok(endpoint)
}
fn chat_completions_url(base_url: &str) -> String {
let base = base_url.trim_end_matches('/');
if base.ends_with("/chat/completions") {
base.to_string()
} else if base.ends_with("/v1") {
format!("{base}/chat/completions")
} else {
format!("{base}/v1/chat/completions")
}
}
fn event_boundary(buffer: &[u8]) -> Option<(usize, usize)> {
buffer
.windows(4)
.position(|window| window == b"\r\n\r\n")
.map(|position| (position, 4))
.or_else(|| {
buffer
.windows(2)
.position(|window| window == b"\n\n")
.map(|position| (position, 2))
})
}
fn in_flight() -> &'static Mutex<HashMap<String, Arc<AtomicBool>>> {
static IN_FLIGHT: OnceLock<Mutex<HashMap<String, Arc<AtomicBool>>>> = OnceLock::new();
IN_FLIGHT.get_or_init(|| Mutex::new(HashMap::new()))
}
fn listeners() -> &'static Mutex<Vec<mpsc::Sender<ChatEvent>>> {
static LISTENERS: OnceLock<Mutex<Vec<mpsc::Sender<ChatEvent>>>> = OnceLock::new();
LISTENERS.get_or_init(|| Mutex::new(Vec::new()))
}
fn emit(options: &ChatSendOptions, event: ChatEvent) {
if let Some(handler) = options.event_handler.as_deref() {
handler(event.clone());
}
listeners()
.lock()
.expect("chat listeners lock")
.retain(|sender| sender.send(event.clone()).is_ok());
}
struct InFlightGuard(String);
impl Drop for InFlightGuard {
fn drop(&mut self) {
in_flight()
.lock()
.expect("chat cancellation lock")
.remove(&self.0);
}
}

861
crates/bds-core/src/engine/chat_tools.rs Normal file
View File

@@ -0,0 +1,861 @@
use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::path::Path;
use base64::Engine as _;
use diesel::prelude::*;
use serde_json::{Value, json};
use crate::db::DbConnection as Connection;
use crate::db::queries::{media, post, post_link, post_media, script, template};
use crate::db::schema::ai_models;
use crate::engine::{EngineError, EngineResult};
use crate::util::frontmatter::read_post_file;
pub fn model_supports_tools(conn: &Connection, model: &str) -> EngineResult<bool> {
let catalog_value = conn.with(|connection| {
ai_models::table
.filter(ai_models::model_id.eq(model))
.select(ai_models::tool_call)
.first::<i32>(connection)
.optional()
})?;
Ok(catalog_value.map_or_else(
|| {
let model = model.to_ascii_lowercase();
model.contains("gpt")
|| model.contains("claude")
|| model.contains("tool")
|| model.contains("qwen")
|| model.contains("mistral")
},
|value| value != 0,
))
}
pub fn system_prompt(conn: &Connection, project_id: &str) -> EngineResult<String> {
let posts = post::list_posts_by_project(conn, project_id)?;
let media_count = media::count_media_by_project(conn, project_id)?;
let tags = posts
.iter()
.flat_map(|post| post.tags.iter())
.collect::<BTreeSet<_>>()
.len();
let categories = posts
.iter()
.flat_map(|post| post.categories.iter())
.collect::<BTreeSet<_>>()
.len();
let configured = crate::engine::settings::get(conn, "ai.system_prompt")?.unwrap_or_default();
let contract = format!(
"You are the conversational assistant for this blog project. Use tools when facts from the project are needed. Never invent project content or identifiers. There are {} posts, {media_count} media items, {tags} tags, and {categories} categories. Keep answers concise and use GitHub-flavored Markdown when useful.",
posts.len()
);
Ok(if configured.trim().is_empty() {
contract
} else {
format!("{}\n\n{contract}", configured.trim())
})
}
pub fn tool_specs() -> Vec<Value> {
vec![
spec(
"get_blog_stats",
"Return aggregate project statistics.",
json!({}),
),
spec(
"check_term",
"Check whether a term is used as a tag or category.",
json!({"term": {"type": "string"}}),
),
spec(
"search_posts",
"Search post titles, slugs, excerpts, bodies, tags, and categories.",
json!({"query": {"type": "string"}, "language": {"type": "string"}, "limit": {"type": "integer"}}),
),
spec(
"read_post",
"Read one post by id.",
json!({"post_id": {"type": "string"}}),
),
spec(
"read_post_by_slug",
"Read one post by slug.",
json!({"slug": {"type": "string"}}),
),
spec(
"list_posts",
"List posts, optionally filtering by status, tag, or category.",
json!({"status": {"type": "string"}, "tag": {"type": "string"}, "category": {"type": "string"}, "limit": {"type": "integer"}}),
),
spec(
"count_posts",
"Count posts and optionally group by status, tag, or category.",
json!({"group_by": {"type": "string", "enum": ["status", "tag", "category"]}}),
),
spec(
"update_post_metadata",
"Update title, excerpt, tags, or categories on a post.",
json!({"post_id": {"type": "string"}, "title": {"type": "string"}, "excerpt": {"type": ["string", "null"]}, "tags": {"type": "array", "items": {"type": "string"}}, "categories": {"type": "array", "items": {"type": "string"}}}),
),
spec(
"list_media",
"List project media.",
json!({"limit": {"type": "integer"}}),
),
spec(
"get_media",
"Get one media item.",
json!({"media_id": {"type": "string"}}),
),
spec(
"view_image",
"Return a local image thumbnail as a data URL for visual inspection.",
json!({"media_id": {"type": "string"}, "size": {"type": "string", "enum": ["small", "medium", "large"]}}),
),
spec(
"update_media_metadata",
"Update title, alt text, caption, or tags on media.",
json!({"media_id": {"type": "string"}, "title": {"type": ["string", "null"]}, "alt": {"type": ["string", "null"]}, "caption": {"type": ["string", "null"]}, "tags": {"type": "array", "items": {"type": "string"}}}),
),
spec("list_tags", "List tags and usage counts.", json!({})),
spec(
"list_categories",
"List categories and usage counts.",
json!({}),
),
spec(
"get_post_backlinks",
"Get posts that link to a post.",
json!({"post_id": {"type": "string"}}),
),
spec(
"get_post_outlinks",
"Get posts linked from a post.",
json!({"post_id": {"type": "string"}}),
),
spec(
"get_post_media",
"Get media linked to a post.",
json!({"post_id": {"type": "string"}}),
),
spec(
"get_media_posts",
"Get posts that use a media item.",
json!({"media_id": {"type": "string"}}),
),
spec("list_templates", "List templates.", json!({})),
spec(
"read_template",
"Read a template by id.",
json!({"template_id": {"type": "string"}}),
),
spec("list_scripts", "List scripts.", json!({})),
spec(
"read_script",
"Read a script by id.",
json!({"script_id": {"type": "string"}}),
),
spec(
"navigate",
"Open a project area or entity in the application.",
json!({
"action": {"type": "string", "enum": ["open_post", "open_media", "open_settings", "open_chat", "switch_view", "toggle_sidebar", "toggle_panel", "toggle_assistant_sidebar"]},
"destination": {"type": "string", "enum": ["posts", "pages", "media", "templates", "scripts", "tags", "chat", "import", "git", "settings"]},
"entity_id": {"type": "string"},
"value": {"type": "string"}
}),
),
]
}
pub fn execute(
conn: &Connection,
data_dir: &Path,
project_id: &str,
name: &str,
arguments: &Value,
) -> EngineResult<Value> {
match name {
"blog_stats" | "get_blog_stats" => blog_stats(conn, project_id),
"check_term" => check_term(conn, project_id, required_str(arguments, "term")?),
"search_posts" => search_posts(conn, project_id, arguments),
"read_post" => {
let item = post::get_post_by_id(conn, required_id(arguments, "post_id", "postId")?)?;
ensure_project(&item.project_id, project_id)?;
post_detail(data_dir, item)
}
"read_post_by_slug" => {
let item = post::get_post_by_project_and_slug(
conn,
project_id,
required_str(arguments, "slug")?,
)?;
post_detail(data_dir, item)
}
"list_posts" => list_posts(conn, project_id, arguments),
"count_posts" => count_posts(conn, project_id, arguments),
"update_post_metadata" => update_post_metadata(conn, data_dir, project_id, arguments),
"list_media" => list_media(conn, project_id, arguments),
"get_media" => {
let item =
media::get_media_by_id(conn, required_id(arguments, "media_id", "mediaId")?)?;
ensure_project(&item.project_id, project_id)?;
Ok(json!({"success": true, "media": item}))
}
"view_image" => view_image(conn, data_dir, project_id, arguments),
"update_media_metadata" => update_media_metadata(conn, data_dir, project_id, arguments),
"list_tags" => counted_terms(conn, project_id, true),
"list_categories" => counted_terms(conn, project_id, false),
"get_post_backlinks" => post_links(conn, project_id, arguments, true),
"get_post_outlinks" => post_links(conn, project_id, arguments, false),
"get_post_media" => linked_media(conn, project_id, arguments),
"get_media_posts" => linked_posts(conn, project_id, arguments),
"list_templates" => Ok(json!({
"templates": template::list_templates_by_project(conn, project_id)?,
})),
"read_template" => {
let item = template::get_template_by_id(
conn,
required_id(arguments, "template_id", "templateId")?,
)?;
ensure_project(&item.project_id, project_id)?;
Ok(json!({"success": true, "template": item}))
}
"list_scripts" => Ok(json!({
"scripts": script::list_scripts_by_project(conn, project_id)?,
})),
"read_script" => {
let item =
script::get_script_by_id(conn, required_id(arguments, "script_id", "scriptId")?)?;
ensure_project(&item.project_id, project_id)?;
Ok(json!({"success": true, "script": item}))
}
"navigate" => navigate(arguments),
_ => Ok(json!({"success": false, "error": "unknown_tool", "name": name})),
}
}
fn blog_stats(conn: &Connection, project_id: &str) -> EngineResult<Value> {
let posts = post::list_posts_by_project(conn, project_id)?;
let media_count = media::count_media_by_project(conn, project_id)?;
let templates = template::list_templates_by_project(conn, project_id)?.len();
let scripts = script::list_scripts_by_project(conn, project_id)?.len();
let tags = posts
.iter()
.flat_map(|item| &item.tags)
.collect::<BTreeSet<_>>()
.len();
let categories = posts
.iter()
.flat_map(|item| &item.categories)
.collect::<BTreeSet<_>>()
.len();
Ok(json!({
"posts": posts.len(), "media": media_count, "templates": templates,
"scripts": scripts, "tags": tags, "categories": categories,
}))
}
fn check_term(conn: &Connection, project_id: &str, term: &str) -> EngineResult<Value> {
let term = term.to_lowercase();
let posts = post::list_posts_by_project(conn, project_id)?;
let tag_count = posts
.iter()
.filter(|item| item.tags.iter().any(|value| value.to_lowercase() == term))
.count();
let category_count = posts
.iter()
.filter(|item| {
item.categories
.iter()
.any(|value| value.to_lowercase() == term)
})
.count();
Ok(json!({"term": term, "tag_posts": tag_count, "category_posts": category_count}))
}
fn search_posts(conn: &Connection, project_id: &str, arguments: &Value) -> EngineResult<Value> {
let query = required_str(arguments, "query")?;
let language = arguments
.get("language")
.and_then(Value::as_str)
.unwrap_or("en");
let limit = limit(arguments);
let mut matches = Vec::new();
for id in crate::db::fts::search_posts(conn, query, language)? {
if let Ok(item) = post::get_post_by_id(conn, &id)
&& item.project_id == project_id
{
matches.push(post_summary(&item));
if matches.len() == limit {
break;
}
}
}
Ok(json!({"posts": matches, "count": matches.len()}))
}
fn list_posts(conn: &Connection, project_id: &str, arguments: &Value) -> EngineResult<Value> {
let status = arguments.get("status").and_then(Value::as_str);
let tag = arguments.get("tag").and_then(Value::as_str);
let category = arguments.get("category").and_then(Value::as_str);
let items = post::list_posts_by_project(conn, project_id)?
.into_iter()
.filter(|item| status.is_none_or(|value| item.status.as_str() == value))
.filter(|item| tag.is_none_or(|value| contains_case_insensitive(&item.tags, value)))
.filter(|item| {
category.is_none_or(|value| contains_case_insensitive(&item.categories, value))
})
.take(limit(arguments))
.map(|item| post_summary(&item))
.collect::<Vec<_>>();
Ok(json!({"posts": items, "count": items.len()}))
}
fn count_posts(conn: &Connection, project_id: &str, arguments: &Value) -> EngineResult<Value> {
let items = post::list_posts_by_project(conn, project_id)?;
let Some(group_by) = arguments
.get("group_by")
.or_else(|| arguments.get("groupBy"))
.and_then(Value::as_str)
else {
return Ok(json!({"total_posts": items.len()}));
};
let mut groups = BTreeMap::<String, usize>::new();
for item in &items {
let values: Vec<String> = match group_by {
"status" => vec![item.status.as_str().to_string()],
"tag" => item.tags.clone(),
"category" => item.categories.clone(),
_ => {
return Err(EngineError::Validation(format!(
"unsupported post grouping: {group_by}"
)));
}
};
for value in values {
*groups.entry(value).or_default() += 1;
}
}
Ok(json!({"total_posts": items.len(), "group_by": group_by, "groups": groups}))
}
fn update_post_metadata(
conn: &Connection,
data_dir: &Path,
project_id: &str,
arguments: &Value,
) -> EngineResult<Value> {
let id = required_id(arguments, "post_id", "postId")?;
let existing = post::get_post_by_id(conn, id)?;
ensure_project(&existing.project_id, project_id)?;
if !["title", "excerpt", "tags", "categories"]
.iter()
.any(|key| arguments.get(key).is_some())
{
return Err(EngineError::Validation(
"no post metadata updates provided".to_string(),
));
}
let excerpt = optional_nullable_str(arguments, "excerpt")?;
let item = crate::engine::post::update_post(
conn,
data_dir,
id,
arguments.get("title").and_then(Value::as_str),
None,
excerpt,
None,
optional_string_array(arguments, "tags")?,
optional_string_array(arguments, "categories")?,
None,
None,
None,
None,
)?;
Ok(json!({"success": true, "post": post_summary(&item)}))
}
fn list_media(conn: &Connection, project_id: &str, arguments: &Value) -> EngineResult<Value> {
let items = media::list_media_by_project(conn, project_id)?
.into_iter()
.take(limit(arguments))
.collect::<Vec<_>>();
Ok(json!({"media": items, "count": items.len()}))
}
fn update_media_metadata(
conn: &Connection,
data_dir: &Path,
project_id: &str,
arguments: &Value,
) -> EngineResult<Value> {
let id = required_id(arguments, "media_id", "mediaId")?;
let existing = media::get_media_by_id(conn, id)?;
ensure_project(&existing.project_id, project_id)?;
if !["title", "alt", "caption", "tags"]
.iter()
.any(|key| arguments.get(key).is_some())
{
return Err(EngineError::Validation(
"no media metadata updates provided".to_string(),
));
}
let item = crate::engine::media::update_media(
conn,
data_dir,
id,
optional_nullable_str(arguments, "title")?,
optional_nullable_str(arguments, "alt")?,
optional_nullable_str(arguments, "caption")?,
None,
None,
optional_string_array(arguments, "tags")?,
)?;
Ok(json!({"success": true, "media": item}))
}
fn view_image(
conn: &Connection,
data_dir: &Path,
project_id: &str,
arguments: &Value,
) -> EngineResult<Value> {
let item = media::get_media_by_id(conn, required_id(arguments, "media_id", "mediaId")?)?;
ensure_project(&item.project_id, project_id)?;
if !item.mime_type.starts_with("image/") {
return Ok(json!({"success": false, "error": "not_image", "mime_type": item.mime_type}));
}
let size = arguments
.get("size")
.and_then(Value::as_str)
.unwrap_or("medium");
if !["small", "medium", "large"].contains(&size) {
return Err(EngineError::Validation(format!(
"unsupported thumbnail size: {size}"
)));
}
let path = data_dir.join(crate::util::thumbnail_path(&item.id, size, "webp"));
if !path.is_file() {
return Ok(json!({"success": false, "error": "thumbnail_not_available"}));
}
let encoded = base64::engine::general_purpose::STANDARD.encode(fs::read(path)?);
Ok(json!({
"success": true,
"media": item,
"data_url": format!("data:image/webp;base64,{encoded}"),
}))
}
fn post_links(
conn: &Connection,
project_id: &str,
arguments: &Value,
incoming: bool,
) -> EngineResult<Value> {
let id = required_id(arguments, "post_id", "postId")?;
let source = post::get_post_by_id(conn, id)?;
ensure_project(&source.project_id, project_id)?;
let links = if incoming {
post_link::list_links_by_target(conn, id)?
} else {
post_link::list_links_by_source(conn, id)?
};
let mut items = Vec::with_capacity(links.len());
for link in links {
let linked_id = if incoming {
&link.source_post_id
} else {
&link.target_post_id
};
let linked = post::get_post_by_id(conn, linked_id)?;
ensure_project(&linked.project_id, project_id)?;
items.push(json!({
"post": post_summary(&linked),
"link_text": link.link_text,
}));
}
if incoming {
Ok(json!({"success": true, "post_id": id, "linked_by": items}))
} else {
Ok(json!({"success": true, "post_id": id, "links_to": items}))
}
}
fn linked_media(conn: &Connection, project_id: &str, arguments: &Value) -> EngineResult<Value> {
let id = required_id(arguments, "post_id", "postId")?;
let item = post::get_post_by_id(conn, id)?;
ensure_project(&item.project_id, project_id)?;
let mut items = Vec::new();
for link in post_media::list_post_media_by_post(conn, id)? {
ensure_project(&link.project_id, project_id)?;
let item = media::get_media_by_id(conn, &link.media_id)?;
ensure_project(&item.project_id, project_id)?;
items.push(json!({"media": item, "sort_order": link.sort_order}));
}
Ok(json!({"success": true, "post_id": id, "media": items}))
}
fn linked_posts(conn: &Connection, project_id: &str, arguments: &Value) -> EngineResult<Value> {
let id = required_id(arguments, "media_id", "mediaId")?;
let item = media::get_media_by_id(conn, id)?;
ensure_project(&item.project_id, project_id)?;
let mut items = Vec::new();
for link in post_media::list_post_media_by_media(conn, id)? {
ensure_project(&link.project_id, project_id)?;
let item = post::get_post_by_id(conn, &link.post_id)?;
ensure_project(&item.project_id, project_id)?;
items.push(json!({"post": post_summary(&item), "sort_order": link.sort_order}));
}
Ok(json!({"success": true, "media_id": id, "posts": items}))
}
fn counted_terms(conn: &Connection, project_id: &str, tags: bool) -> EngineResult<Value> {
let mut counts = BTreeMap::<String, usize>::new();
for item in post::list_posts_by_project(conn, project_id)? {
for term in if tags { &item.tags } else { &item.categories } {
*counts.entry(term.clone()).or_default() += 1;
}
}
let values = counts
.into_iter()
.map(|(name, count)| json!({"name": name, "count": count}))
.collect::<Vec<_>>();
Ok(if tags {
json!({"tags": values, "count": values.len()})
} else {
json!({"categories": values, "count": values.len()})
})
}
fn navigate(arguments: &Value) -> EngineResult<Value> {
let action = arguments.get("action").and_then(Value::as_str);
let value = arguments
.get("value")
.or_else(|| arguments.get("entity_id"))
.or_else(|| arguments.get("entityId"))
.and_then(Value::as_str);
let (destination, entity_id) = match action {
Some("open_post" | "openPost") => ("posts", required_navigation_value(value, "post")?),
Some("open_media" | "openMedia") => ("media", required_navigation_value(value, "media")?),
Some("open_chat" | "openChat") => ("chat", required_navigation_value(value, "chat")?),
Some("open_settings" | "openSettings") => ("settings", None),
Some("switch_view" | "switchView") => {
(required_navigation_value(value, "view")?.unwrap(), None)
}
Some("toggle_sidebar" | "toggleSidebar") => ("toggle_sidebar", None),
Some("toggle_panel" | "togglePanel") => ("toggle_panel", None),
Some("toggle_assistant_sidebar" | "toggleAssistantSidebar") => {
("toggle_assistant_sidebar", None)
}
Some(action) => {
return Err(EngineError::Validation(format!(
"unsupported navigation action: {action}"
)));
}
None => (required_str(arguments, "destination")?, value),
};
if ![
"posts",
"pages",
"media",
"templates",
"scripts",
"tags",
"chat",
"import",
"git",
"settings",
"toggle_sidebar",
"toggle_panel",
"toggle_assistant_sidebar",
]
.contains(&destination)
{
return Err(EngineError::Validation(format!(
"unsupported navigation destination: {destination}"
)));
}
Ok(json!({
"success": true,
"navigation": {
"destination": destination,
"entity_id": entity_id,
}
}))
}
fn post_detail(data_dir: &Path, item: crate::model::Post) -> EngineResult<Value> {
let body = post_body(data_dir, &item)?;
Ok(json!({"success": true, "post": item, "body": body}))
}
fn post_body(data_dir: &Path, item: &crate::model::Post) -> EngineResult<String> {
if let Some(content) = item.content.as_deref() {
return Ok(content.to_string());
}
if item.file_path.is_empty() {
return Ok(String::new());
}
let raw = fs::read_to_string(data_dir.join(&item.file_path))?;
let (_, body) = read_post_file(&raw)
.map_err(|error| EngineError::Parse(format!("invalid post file: {error}")))?;
Ok(body)
}
fn post_summary(item: &crate::model::Post) -> Value {
json!({
"id": item.id, "title": item.title, "slug": item.slug,
"excerpt": item.excerpt, "status": item.status, "tags": item.tags,
"categories": item.categories, "created_at": item.created_at,
"updated_at": item.updated_at,
})
}
fn spec(name: &str, description: &str, properties: Value) -> Value {
let required = properties
.as_object()
.into_iter()
.flat_map(|values| values.iter())
.filter(|(_, schema)| !schema.get("type").is_some_and(Value::is_array))
.filter(|(name, _)| {
matches!(
name.as_str(),
"term"
| "query"
| "post_id"
| "slug"
| "media_id"
| "template_id"
| "script_id"
| "destination"
)
})
.map(|(name, _)| Value::String(name.clone()))
.collect::<Vec<_>>();
let required = if name == "navigate" {
Vec::new()
} else {
required
};
json!({
"type": "function",
"function": {
"name": name,
"description": description,
"parameters": {
"type": "object",
"properties": properties,
"required": required,
"additionalProperties": false,
}
}
})
}
fn required_navigation_value<'a>(
value: Option<&'a str>,
target: &str,
) -> EngineResult<Option<&'a str>> {
value
.filter(|value| !value.trim().is_empty())
.map(Some)
.ok_or_else(|| {
EngineError::Validation(format!("navigation {target} identifier is required"))
})
}
fn required_str<'a>(arguments: &'a Value, key: &str) -> EngineResult<&'a str> {
arguments
.get(key)
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.ok_or_else(|| EngineError::Validation(format!("tool argument {key} is required")))
}
fn required_id<'a>(arguments: &'a Value, snake: &str, camel: &str) -> EngineResult<&'a str> {
arguments
.get(snake)
.or_else(|| arguments.get(camel))
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.ok_or_else(|| EngineError::Validation(format!("tool argument {snake} is required")))
}
fn ensure_project(actual: &str, expected: &str) -> EngineResult<()> {
if actual == expected {
Ok(())
} else {
Err(EngineError::NotFound("project entity".to_string()))
}
}
fn limit(arguments: &Value) -> usize {
arguments
.get("limit")
.and_then(Value::as_u64)
.unwrap_or(25)
.clamp(1, 100) as usize
}
fn contains_case_insensitive(values: &[String], needle: &str) -> bool {
values
.iter()
.any(|value| value.eq_ignore_ascii_case(needle))
}
fn optional_string_array(arguments: &Value, key: &str) -> EngineResult<Option<Vec<String>>> {
let Some(value) = arguments.get(key) else {
return Ok(None);
};
let values = value
.as_array()
.ok_or_else(|| EngineError::Validation(format!("tool argument {key} must be an array")))?
.iter()
.map(|value| {
value.as_str().map(str::to_string).ok_or_else(|| {
EngineError::Validation(format!("tool argument {key} must contain strings"))
})
})
.collect::<EngineResult<Vec<_>>>()?;
Ok(Some(values))
}
fn optional_nullable_str<'a>(
arguments: &'a Value,
key: &str,
) -> EngineResult<Option<Option<&'a str>>> {
match arguments.get(key) {
None => Ok(None),
Some(Value::Null) => Ok(Some(None)),
Some(Value::String(value)) => Ok(Some(Some(value))),
Some(_) => Err(EngineError::Validation(format!(
"tool argument {key} must be a string or null"
))),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::Database;
use crate::db::queries::media::{insert_media, make_test_media};
use crate::db::queries::post::{insert_post, make_test_post};
use crate::db::queries::post_link::insert_post_link;
use crate::db::queries::post_media::link_media;
use crate::db::queries::project::{insert_project, make_test_project};
use crate::model::{PostLink, PostMedia};
fn setup() -> Database {
let db = Database::open_in_memory().unwrap();
db.migrate().unwrap();
insert_project(db.conn(), &make_test_project("p1", "blog")).unwrap();
insert_post(db.conn(), &make_test_post("source", "p1", "source")).unwrap();
insert_post(db.conn(), &make_test_post("target", "p1", "target")).unwrap();
insert_media(db.conn(), &make_test_media("media1", "p1")).unwrap();
insert_post_link(
db.conn(),
&PostLink {
id: "link1".into(),
source_post_id: "source".into(),
target_post_id: "target".into(),
link_text: Some("read next".into()),
created_at: 1,
},
)
.unwrap();
link_media(
db.conn(),
&PostMedia {
id: "post-media1".into(),
project_id: "p1".into(),
post_id: "source".into(),
media_id: "media1".into(),
sort_order: 3,
created_at: 1,
},
)
.unwrap();
db
}
#[test]
fn relationship_tools_return_project_entities() {
let db = setup();
let dir = tempfile::tempdir().unwrap();
let outlinks = execute(
db.conn(),
dir.path(),
"p1",
"get_post_outlinks",
&json!({"post_id": "source"}),
)
.unwrap();
assert_eq!(outlinks["links_to"][0]["post"]["id"], "target");
assert_eq!(outlinks["links_to"][0]["link_text"], "read next");
let backlinks = execute(
db.conn(),
dir.path(),
"p1",
"get_post_backlinks",
&json!({"post_id": "target"}),
)
.unwrap();
assert_eq!(backlinks["linked_by"][0]["post"]["id"], "source");
let post_media = execute(
db.conn(),
dir.path(),
"p1",
"get_post_media",
&json!({"post_id": "source"}),
)
.unwrap();
assert_eq!(post_media["media"][0]["media"]["id"], "media1");
assert_eq!(post_media["media"][0]["sort_order"], 3);
let media_posts = execute(
db.conn(),
dir.path(),
"p1",
"get_media_posts",
&json!({"media_id": "media1"}),
)
.unwrap();
assert_eq!(media_posts["posts"][0]["post"]["id"], "source");
}
#[test]
fn view_image_is_bounded_to_generated_image_thumbnails() {
let db = setup();
let dir = tempfile::tempdir().unwrap();
let relative = crate::util::thumbnail_path("media1", "medium", "webp");
let path = dir.path().join(relative);
fs::create_dir_all(path.parent().unwrap()).unwrap();
fs::write(&path, b"thumbnail").unwrap();
let result = execute(
db.conn(),
dir.path(),
"p1",
"view_image",
&json!({"media_id": "media1", "size": "medium"}),
)
.unwrap();
assert_eq!(result["success"], true);
assert_eq!(result["data_url"], "data:image/webp;base64,dGh1bWJuYWls");
let invalid = execute(
db.conn(),
dir.path(),
"p1",
"view_image",
&json!({"media_id": "media1", "size": "original"}),
);
assert!(matches!(invalid, Err(EngineError::Validation(_))));
}
}

4
crates/bds-core/src/engine/git.rs
View File

@@ -1635,7 +1635,7 @@ mod tests {
dir.path(),
&executable,
Duration::from_secs(1),
Duration::from_secs(1),
Duration::from_secs(3),
);
let error = engine.fetch(|| false, |_| {}).unwrap_err();
assert!(
@@ -1658,7 +1658,7 @@ mod tests {
dir.path(),
executable,
Duration::from_secs(1),
Duration::from_secs(1),
Duration::from_secs(3),
);
let error = engine
.push(

2
crates/bds-core/src/engine/mod.rs
View File

@@ -2,6 +2,8 @@ pub mod ai;
pub mod auto_translation;
pub mod blogmark;
pub mod calendar;
pub mod chat;
mod chat_tools;
pub mod cli_launcher;
pub mod cli_sync;
pub mod domain_events;

109
crates/bds-core/src/model/chat.rs Normal file
View File

@@ -0,0 +1,109 @@
use serde::{Deserialize, Serialize};
#[derive(
Debug,
Clone,
Copy,
PartialEq,
Eq,
Serialize,
Deserialize,
diesel::AsExpression,
diesel::FromSqlRow,
)]
#[diesel(sql_type = diesel::sql_types::Text)]
#[serde(rename_all = "lowercase")]
pub enum ChatRole {
System,
User,
Assistant,
Tool,
}
impl ChatRole {
pub const fn as_str(self) -> &'static str {
match self {
Self::System => "system",
Self::User => "user",
Self::Assistant => "assistant",
Self::Tool => "tool",
}
}
}
impl std::str::FromStr for ChatRole {
type Err = String;
fn from_str(value: &str) -> Result<Self, Self::Err> {
match value {
"system" => Ok(Self::System),
"user" => Ok(Self::User),
"assistant" => Ok(Self::Assistant),
"tool" => Ok(Self::Tool),
_ => Err(format!("invalid chat role: {value}")),
}
}
}
#[derive(
Debug, Clone, PartialEq, Eq, Serialize, Deserialize, diesel::Queryable, diesel::Selectable,
)]
#[diesel(
table_name = crate::db::schema::chat_conversations,
check_for_backend(diesel::sqlite::Sqlite)
)]
pub struct ChatConversation {
pub id: String,
pub title: String,
pub model: Option<String>,
pub copilot_session_id: Option<String>,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Debug, diesel::Insertable)]
#[diesel(table_name = crate::db::schema::chat_conversations)]
pub struct NewChatConversation<'a> {
pub id: &'a str,
pub title: &'a str,
pub model: Option<&'a str>,
pub copilot_session_id: Option<&'a str>,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(
Debug, Clone, PartialEq, Eq, Serialize, Deserialize, diesel::Queryable, diesel::Selectable,
)]
#[diesel(
table_name = crate::db::schema::chat_messages,
check_for_backend(diesel::sqlite::Sqlite)
)]
pub struct ChatMessage {
pub id: i32,
pub conversation_id: String,
pub role: ChatRole,
pub content: Option<String>,
pub tool_call_id: Option<String>,
pub tool_calls: Option<String>,
pub created_at: i64,
pub cache_read_tokens: Option<i32>,
pub cache_write_tokens: Option<i32>,
pub token_usage_input: Option<i32>,
pub token_usage_output: Option<i32>,
}
#[derive(Debug, diesel::Insertable)]
#[diesel(table_name = crate::db::schema::chat_messages)]
pub struct NewChatMessage<'a> {
pub conversation_id: &'a str,
pub role: ChatRole,
pub content: Option<&'a str>,
pub tool_call_id: Option<&'a str>,
pub tool_calls: Option<&'a str>,
pub created_at: i64,
pub cache_read_tokens: Option<i32>,
pub cache_write_tokens: Option<i32>,
pub token_usage_input: Option<i32>,
pub token_usage_output: Option<i32>,
}

2
crates/bds-core/src/model/mod.rs
View File

@@ -1,3 +1,4 @@
mod chat;
mod event;
mod generation;
mod import;
@@ -10,6 +11,7 @@ mod script;
mod tag;
mod template;
pub use chat::{ChatConversation, ChatMessage, ChatRole, NewChatConversation, NewChatMessage};
pub use event::DomainEvent;
pub use generation::{
DbNotification, DomainEntity, GeneratedFileHash, NotificationAction, NotificationEntity,

570
crates/bds-core/tests/chat.rs Normal file
View File

@@ -0,0 +1,570 @@
use bds_core::db::Database;
use bds_core::engine::ai::{AiEndpointConfig, AiEndpointKind};
use bds_core::engine::{chat, project};
use bds_core::model::ChatRole;
use std::io::{Read, Write};
use std::net::TcpListener;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
fn setup() -> (tempfile::TempDir, Database, String, std::path::PathBuf) {
let root = tempfile::tempdir().unwrap();
let database_path = root.path().join("bds.db");
let data_dir = root.path().join("project");
std::fs::create_dir_all(&data_dir).unwrap();
let db = Database::open(&database_path).unwrap();
db.migrate().unwrap();
let project = project::create_project(db.conn(), "Chat Test", data_dir.to_str()).unwrap();
project::set_active_project(db.conn(), &project.id).unwrap();
(root, db, project.id, data_dir)
}
#[test]
fn conversation_repository_round_trips_rename_model_messages_and_delete() {
let (_root, db, _project_id, _data_dir) = setup();
let conversation = chat::create_conversation(db.conn(), Some("tool-model")).unwrap();
assert_eq!(conversation.title, "Chat with tool-model");
assert_eq!(conversation.model.as_deref(), Some("tool-model"));
let renamed = chat::rename_conversation(db.conn(), &conversation.id, "Rust notes").unwrap();
assert_eq!(renamed.title, "Rust notes");
chat::set_conversation_model(db.conn(), &conversation.id, "other-model").unwrap();
chat::insert_message(
db.conn(),
&conversation.id,
ChatRole::User,
Some("Hello"),
None,
None,
Default::default(),
)
.unwrap();
assert_eq!(chat::list_conversations(db.conn()).unwrap().len(), 1);
assert_eq!(
chat::list_messages(db.conn(), &conversation.id)
.unwrap()
.len(),
1
);
assert_eq!(
chat::list_messages(db.conn(), &conversation.id).unwrap()[0].role,
ChatRole::User
);
chat::delete_conversation(db.conn(), &conversation.id).unwrap();
assert!(chat::list_conversations(db.conn()).unwrap().is_empty());
assert!(
chat::list_messages(db.conn(), &conversation.id)
.unwrap()
.is_empty()
);
}
#[test]
fn sse_assembler_handles_split_frames_multiple_tools_and_usage() {
let mut assembler = chat::SseAssembler::default();
assembler
.feed(b"data: {\"choices\":[{\"delta\":{\"content\":\"Hel")
.unwrap();
assembler
.feed(b"lo\"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call-1\",\"function\":{\"name\":\"search_posts\",\"arguments\":\"{\\\"query\\\":\"}},{\"index\":1,\"id\":\"call-2\",\"function\":{\"name\":\"get_blog_stats\",\"arguments\":\"{}\"}}]}}]}\n\n")
.unwrap();
assembler
.feed(b"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"rust\\\"}\"}}]}}],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":4,\"prompt_tokens_details\":{\"cached_tokens\":7},\"completion_tokens_details\":{\"cached_tokens\":2}}}\n\ndata: [DONE]\n\n")
.unwrap();
let assembled = assembler.finish().unwrap();
assert_eq!(assembled.content, "Hello");
assert_eq!(assembled.tool_calls.len(), 2);
assert_eq!(assembled.tool_calls[0].name, "search_posts");
assert_eq!(assembled.tool_calls[0].arguments["query"], "rust");
assert_eq!(assembled.usage.input_tokens, Some(11));
assert_eq!(assembled.usage.output_tokens, Some(4));
assert_eq!(assembled.usage.cache_read_tokens, Some(7));
assert_eq!(assembled.usage.cache_write_tokens, Some(2));
}
#[test]
fn unavailable_endpoint_refuses_chat_without_mutating_transcript() {
let (_root, db, project_id, data_dir) = setup();
let conversation = chat::create_conversation(db.conn(), None).unwrap();
let result = chat::send_chat_message(
db.conn(),
&data_dir,
&project_id,
false,
&conversation.id,
"Hello",
chat::ChatSendOptions::default(),
);
assert!(result.is_err());
assert!(
chat::list_messages(db.conn(), &conversation.id)
.unwrap()
.is_empty()
);
}
#[test]
fn streamed_turn_persists_content_session_and_all_token_fields() {
let (_root, db, project_id, data_dir) = setup();
let conversation = chat::create_conversation(db.conn(), Some("plain-model")).unwrap();
let (url, server) = serve(vec![MockResponse::delayed_sse(vec![
"data: {\"session_id\":\"session-7\",\"choices\":[{\"delta\":{\"content\":\"Hi \"}}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"content\":\"there\"}}],\"usage\":{\"prompt_tokens\":12,\"completion_tokens\":3,\"cache_read_tokens\":4,\"cache_write_tokens\":2}}\n\n",
"data: [DONE]\n\n",
])]);
let snapshots = Arc::new(std::sync::Mutex::new(Vec::new()));
let captured = Arc::clone(&snapshots);
let result = chat::send_chat_message(
db.conn(),
&data_dir,
&project_id,
false,
&conversation.id,
"Hello",
options(url, move |event| {
if let chat::ChatEvent::Content { content, .. } = event {
captured.lock().unwrap().push(content);
}
}),
)
.unwrap();
server.join().unwrap();
assert_eq!(result.content, "Hi there");
assert_eq!(&*snapshots.lock().unwrap(), &["Hi ", "Hi there"]);
let messages = chat::list_messages(db.conn(), &conversation.id).unwrap();
assert_eq!(messages.len(), 2);
assert_eq!(messages[1].content.as_deref(), Some("Hi there"));
assert_eq!(messages[1].token_usage_input, Some(12));
assert_eq!(messages[1].token_usage_output, Some(3));
assert_eq!(messages[1].cache_read_tokens, Some(4));
assert_eq!(messages[1].cache_write_tokens, Some(2));
assert_eq!(
chat::get_conversation(db.conn(), &conversation.id)
.unwrap()
.copilot_session_id
.as_deref(),
Some("session-7")
);
}
#[test]
fn tool_loop_persists_valid_assistant_tool_pair_before_final_answer() {
let (_root, db, project_id, data_dir) = setup();
let conversation = chat::create_conversation(db.conn(), Some("tool-model")).unwrap();
let (url, server) = serve(vec![
MockResponse::sse(vec![
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"stats-1\",\"function\":{\"name\":\"get_blog_stats\",\"arguments\":\"{}\"}}]}}]}\n\n",
"data: [DONE]\n\n",
]),
MockResponse::sse(vec![
"data: {\"choices\":[{\"delta\":{\"content\":\"There are no posts.\"}}]}\n\n",
"data: [DONE]\n\n",
]),
]);
let result = chat::send_chat_message(
db.conn(),
&data_dir,
&project_id,
false,
&conversation.id,
"How many posts?",
options(url, |_| {}),
)
.unwrap();
server.join().unwrap();
assert_eq!(result.content, "There are no posts.");
let messages = chat::list_messages(db.conn(), &conversation.id).unwrap();
assert_eq!(
messages.iter().map(|item| item.role).collect::<Vec<_>>(),
vec![
ChatRole::User,
ChatRole::Assistant,
ChatRole::Tool,
ChatRole::Assistant,
]
);
assert!(
messages[1]
.tool_calls
.as_deref()
.unwrap()
.contains("stats-1")
);
assert_eq!(messages[2].tool_call_id.as_deref(), Some("stats-1"));
assert!(
messages[2]
.content
.as_deref()
.unwrap()
.contains("\"posts\":0")
);
}
#[test]
fn malformed_stream_and_provider_error_keep_reopenable_user_turn() {
for response in [
MockResponse::sse(vec!["data: {not-json}\n\n"]),
MockResponse::status(500, "application/json", "{\"error\":\"down\"}"),
] {
let (_root, db, project_id, data_dir) = setup();
let conversation = chat::create_conversation(db.conn(), Some("plain-model")).unwrap();
let (url, server) = serve(vec![response]);
assert!(
chat::send_chat_message(
db.conn(),
&data_dir,
&project_id,
false,
&conversation.id,
"Remember this",
options(url, |_| {}),
)
.is_err()
);
server.join().unwrap();
let messages = chat::list_messages(db.conn(), &conversation.id).unwrap();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].role, ChatRole::User);
assert_eq!(messages[0].content.as_deref(), Some("Remember this"));
}
}
#[test]
fn cancellation_persists_only_received_content_and_never_runs_later_work() {
let (_root, db, project_id, data_dir) = setup();
let conversation = chat::create_conversation(db.conn(), Some("plain-model")).unwrap();
let conversation_id = conversation.id.clone();
let (url, server) = serve(vec![MockResponse::delayed_sse(vec![
"data: {\"choices\":[{\"delta\":{\"content\":\"Partial\"}}]}\n\n",
"data: {\"choices\":[{\"delta\":{\"content\":\" ignored\"}}]}\n\n",
"data: [DONE]\n\n",
])]);
let cancel_id = conversation_id.clone();
let result = chat::send_chat_message(
db.conn(),
&data_dir,
&project_id,
false,
&conversation_id,
"Start",
options(url, move |event| {
if matches!(event, chat::ChatEvent::Content { .. }) {
chat::cancel_chat(&cancel_id);
}
}),
)
.unwrap();
server.join().unwrap();
assert!(result.cancelled);
assert_eq!(result.content, "Partial");
let messages = chat::list_messages(db.conn(), &conversation_id).unwrap();
assert_eq!(messages.len(), 2);
assert_eq!(messages[1].content.as_deref(), Some("Partial"));
}
#[test]
fn context_truncation_keeps_system_and_complete_tool_pairs() {
let (_root, db, _project_id, _data_dir) = setup();
let conversation = chat::create_conversation(db.conn(), None).unwrap();
chat::insert_message(
db.conn(),
&conversation.id,
ChatRole::User,
Some(&"old ".repeat(100)),
None,
None,
Default::default(),
)
.unwrap();
chat::insert_message(
db.conn(),
&conversation.id,
ChatRole::Assistant,
Some("old answer"),
None,
None,
Default::default(),
)
.unwrap();
chat::insert_message(
db.conn(),
&conversation.id,
ChatRole::User,
Some("new question"),
None,
None,
Default::default(),
)
.unwrap();
chat::insert_message(db.conn(), &conversation.id, ChatRole::Assistant, None, None, Some("[{\"id\":\"c1\",\"type\":\"function\",\"function\":{\"name\":\"get_blog_stats\",\"arguments\":\"{}\"}}]"), Default::default()).unwrap();
chat::insert_message(
db.conn(),
&conversation.id,
ChatRole::Tool,
Some("{\"posts\":0}"),
Some("c1"),
None,
Default::default(),
)
.unwrap();
let context = chat::build_context(db.conn(), &_project_id, &conversation.id, 100).unwrap();
assert_eq!(context[0]["role"], "system");
assert!(context.iter().any(|item| item["content"] == "new question"));
let assistant_index = context
.iter()
.position(|item| item.get("tool_calls").is_some())
.unwrap();
assert_eq!(context[assistant_index + 1]["tool_call_id"], "c1");
assert!(!context.iter().any(|item| {
item["content"]
.as_str()
.is_some_and(|value| value.starts_with("old "))
}));
}
#[test]
fn airplane_mode_selects_the_local_endpoint_and_persists_its_model() {
let (_root, db, project_id, data_dir) = setup();
let conversation = chat::create_conversation(db.conn(), None).unwrap();
let (url, server) = serve(vec![MockResponse::sse(vec![
"data: {\"choices\":[{\"delta\":{\"content\":\"Local answer\"}}]}\n\n",
"data: [DONE]\n\n",
])]);
bds_core::engine::ai::save_endpoint(
db.conn(),
&AiEndpointConfig {
kind: AiEndpointKind::Airplane,
url,
model: "local-airplane-model".to_string(),
api_key: None,
},
)
.unwrap();
let result = chat::send_chat_message(
db.conn(),
&data_dir,
&project_id,
true,
&conversation.id,
"Stay offline",
Default::default(),
)
.unwrap();
server.join().unwrap();
assert_eq!(result.content, "Local answer");
assert_eq!(
chat::get_conversation(db.conn(), &conversation.id)
.unwrap()
.model
.as_deref(),
Some("local-airplane-model")
);
}
#[test]
fn tool_round_limit_stops_without_leaving_an_unpaired_tool_call() {
let (_root, db, project_id, data_dir) = setup();
let conversation = chat::create_conversation(db.conn(), Some("tool-model")).unwrap();
let tool_frame = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"stats\",\"function\":{\"name\":\"get_blog_stats\",\"arguments\":\"{}\"}}]}}]}\n\ndata: [DONE]\n\n";
let (url, server) = serve(vec![
MockResponse::sse(vec![tool_frame]),
MockResponse::sse(vec![tool_frame]),
]);
let mut send_options = options(url, |_| {});
send_options.max_tool_rounds = 1;
let result = chat::send_chat_message(
db.conn(),
&data_dir,
&project_id,
false,
&conversation.id,
"Loop forever",
send_options,
);
assert!(result.is_err());
server.join().unwrap();
let messages = chat::list_messages(db.conn(), &conversation.id).unwrap();
assert_eq!(messages.len(), 3);
assert_eq!(messages[1].role, ChatRole::Assistant);
assert_eq!(messages[2].role, ChatRole::Tool);
assert_eq!(messages[2].tool_call_id.as_deref(), Some("stats"));
assert!(messages[1].tool_calls.as_deref().unwrap().contains("stats"));
}
#[test]
fn cancellation_before_mutating_tool_execution_records_pairs_without_mutation() {
let (_root, db, project_id, data_dir) = setup();
bds_core::db::fts::ensure_fts_tables(db.conn()).unwrap();
let post = bds_core::engine::post::create_post(
db.conn(),
&data_dir,
&project_id,
"Original",
Some("Body"),
Vec::new(),
Vec::new(),
None,
Some("en"),
None,
)
.unwrap();
let conversation = chat::create_conversation(db.conn(), Some("tool-model")).unwrap();
let tool_frame = format!(
"data: {{\"choices\":[{{\"delta\":{{\"tool_calls\":[{{\"index\":0,\"id\":\"update-1\",\"function\":{{\"name\":\"update_post_metadata\",\"arguments\":\"{{\\\"post_id\\\":\\\"{}\\\",\\\"title\\\":\\\"Changed\\\"}}\"}}}},{{\"index\":1,\"id\":\"stats-2\",\"function\":{{\"name\":\"get_blog_stats\",\"arguments\":\"{{}}\"}}}}]}}}}]}}\n\ndata: [DONE]\n\n",
post.id
);
let (url, server) = serve(vec![MockResponse::sse(vec![tool_frame.as_str()])]);
let cancel_id = conversation.id.clone();
let result = chat::send_chat_message(
db.conn(),
&data_dir,
&project_id,
false,
&conversation.id,
"Change it",
options(url, move |event| {
if matches!(event, chat::ChatEvent::ToolStarted { .. }) {
chat::cancel_chat(&cancel_id);
}
}),
)
.unwrap();
server.join().unwrap();
assert!(result.cancelled);
assert_eq!(
bds_core::db::queries::post::get_post_by_id(db.conn(), &post.id)
.unwrap()
.title,
"Original"
);
let messages = chat::list_messages(db.conn(), &conversation.id).unwrap();
assert_eq!(messages.len(), 4);
assert_eq!(messages[2].tool_call_id.as_deref(), Some("update-1"));
assert_eq!(messages[3].tool_call_id.as_deref(), Some("stats-2"));
assert!(
messages[2]
.content
.as_deref()
.unwrap()
.contains("cancelled")
);
}
fn options(
url: String,
handler: impl Fn(chat::ChatEvent) + Send + Sync + 'static,
) -> chat::ChatSendOptions {
chat::ChatSendOptions {
endpoint: Some(AiEndpointConfig {
kind: AiEndpointKind::Online,
url,
model: "plain-model".to_string(),
api_key: Some("test-key".to_string()),
}),
event_handler: Some(Arc::new(handler)),
..Default::default()
}
}
struct MockResponse {
status: u16,
content_type: &'static str,
chunks: Vec<String>,
delay: bool,
}
impl MockResponse {
fn sse(chunks: Vec<&str>) -> Self {
Self {
status: 200,
content_type: "text/event-stream",
chunks: chunks.into_iter().map(str::to_string).collect(),
delay: false,
}
}
fn delayed_sse(chunks: Vec<&str>) -> Self {
Self {
status: 200,
content_type: "text/event-stream",
chunks: chunks.into_iter().map(str::to_string).collect(),
delay: true,
}
}
fn status(status: u16, content_type: &'static str, body: &str) -> Self {
Self {
status,
content_type,
chunks: vec![body.to_string()],
delay: false,
}
}
}
fn serve(responses: Vec<MockResponse>) -> (String, thread::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let address = listener.local_addr().unwrap();
let server = thread::spawn(move || {
for response in responses {
let (mut socket, _) = listener.accept().unwrap();
read_request(&mut socket);
let reason = if response.status == 200 {
"OK"
} else {
"Error"
};
write!(
socket,
"HTTP/1.1 {} {}\r\nContent-Type: {}\r\nConnection: close\r\n\r\n",
response.status, reason, response.content_type
)
.unwrap();
for chunk in response.chunks {
if socket.write_all(chunk.as_bytes()).is_err() || socket.flush().is_err() {
break;
}
if response.delay {
thread::sleep(Duration::from_millis(60));
}
}
}
});
(format!("http://{address}"), server)
}
fn read_request(socket: &mut std::net::TcpStream) {
socket
.set_read_timeout(Some(Duration::from_secs(2)))
.unwrap();
let mut received = Vec::new();
let mut chunk = [0_u8; 2048];
loop {
let count = socket.read(&mut chunk).unwrap();
received.extend_from_slice(&chunk[..count]);
let Some(headers_end) = received.windows(4).position(|value| value == b"\r\n\r\n") else {
continue;
};
let headers = String::from_utf8_lossy(&received[..headers_end]);
let content_length = headers
.lines()
.find_map(|line| {
line.to_ascii_lowercase()
.strip_prefix("content-length:")
.map(str::trim)
.map(str::to_string)
})
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(0);
if received.len() >= headers_end + 4 + content_length {
break;
}
}
}