Cameron CordesandClaude Opus 5.5
305eda5894
Long insight/persona chats on llama-swap failed with "exceeds the available context size" even after raising Context in the app: - llama-server fixes its window at launch (-c) and ignores num_ctx, but the history budget trusted the requested num_ctx, so raising it kept MORE history than the server could hold. LlmClient::context_window() now reports the served window (llama-swap: /upstream/<model>/props, cached 10 min) and the budget uses min(requested, served). - apply_context_budget could only drop tool-call pairs, so plain Q&A chats never shrank. It now also drops whole earlier exchanges (keeping user/assistant alternation) and, last, the current turn's tool results, never the in-progress question. - The budget only ran once, before the tool loop. ContextFitter now fits every LLM call (on a copy; the stored transcript is untouched) and, on a context-overflow error, shrinks the budget from the server's reported token counts and retries once. GET /insights/models reports each llama-swap model's context_window when it's known without loading the model (cached, or currently running), for the app's Context picker. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
809 lines
32 KiB
Rust
809 lines
32 KiB
Rust
use anyhow::Result;
|
|
use async_trait::async_trait;
|
|
use futures::stream::BoxStream;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Provider-agnostic surface for LLM backends (Ollama, OpenRouter, …).
|
|
///
|
|
/// Impls translate these canonical shapes at the wire boundary: tool-call
|
|
/// arguments stay as `serde_json::Value` in memory and are stringified only
|
|
/// when a provider requires it (OpenAI-compatible APIs do), and `images`
|
|
/// stays as base64 strings here and is rewritten into content-parts where
|
|
/// needed.
|
|
///
|
|
/// Every method that performs an HTTP call starts an `llm.request` span
|
|
/// parented on the caller's ambient `tracing` span, so callers must run
|
|
/// inside their feature span (`#[instrument]` / `.instrument(span)`) — and
|
|
/// carry it across any `tokio::spawn` — for the LLM span to nest under it.
|
|
// First consumer lands in a later PR (OpenRouter impl + hybrid mode routing).
|
|
#[allow(dead_code)]
|
|
#[async_trait]
|
|
pub trait LlmClient: Send + Sync {
|
|
/// Single-shot text generation. Optional system prompt and optional
|
|
/// base64 images (ignored by providers without vision support).
|
|
async fn generate(
|
|
&self,
|
|
prompt: &str,
|
|
system: Option<&str>,
|
|
images: Option<Vec<String>>,
|
|
) -> Result<String>;
|
|
|
|
/// Multi-turn chat with tool definitions. Returns the assistant message
|
|
/// (which may contain tool_calls) plus optional prompt/eval token counts.
|
|
async fn chat_with_tools(
|
|
&self,
|
|
messages: Vec<ChatMessage>,
|
|
tools: Vec<Tool>,
|
|
) -> Result<(ChatMessage, Option<i32>, Option<i32>)>;
|
|
|
|
/// Streaming variant of `chat_with_tools`. The returned stream yields
|
|
/// `TextDelta` items as content is produced, then a single terminal
|
|
/// `Done` carrying the complete assembled message (with tool_calls, if
|
|
/// any) plus token usage counts. Implementations that can't stream may
|
|
/// fall back to calling `chat_with_tools` and emitting the full reply
|
|
/// as one `Done` event.
|
|
async fn chat_with_tools_stream(
|
|
&self,
|
|
messages: Vec<ChatMessage>,
|
|
tools: Vec<Tool>,
|
|
) -> Result<BoxStream<'static, Result<LlmStreamEvent>>>;
|
|
|
|
/// Batch embedding generation. Dimensionality is provider/model specific.
|
|
async fn generate_embeddings(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
|
|
|
|
/// One-shot vision description of an image. Used to convert images into
|
|
/// plain text for the hybrid-mode conversation flow.
|
|
async fn describe_image(&self, image_base64: &str) -> Result<String>;
|
|
|
|
/// Enumerate available models with their capabilities.
|
|
async fn list_models(&self) -> Result<Vec<ModelCapabilities>>;
|
|
|
|
/// Look up capabilities for a single model.
|
|
async fn model_capabilities(&self, model: &str) -> Result<ModelCapabilities>;
|
|
|
|
/// Primary model identifier this client was constructed with.
|
|
fn primary_model(&self) -> &str;
|
|
|
|
/// Context window (tokens) the server actually serves for the primary
|
|
/// model, when the backend can report one that overrides the request's
|
|
/// `num_ctx`. `None` = unknown or request-controlled (Ollama honours
|
|
/// `num_ctx` per request); callers fall back to the requested value.
|
|
async fn context_window(&self) -> Option<usize> {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Events emitted by streaming `chat_with_tools_stream`. A stream is a
|
|
/// sequence of zero or more `TextDelta` events followed by exactly one
|
|
/// `Done`. Callers should treat `Done` as terminal — further items (if any
|
|
/// slip through due to upstream misbehavior) are safe to ignore.
|
|
#[derive(Debug, Clone)]
|
|
pub enum LlmStreamEvent {
|
|
/// Incremental content token(s) from the model. Concatenate in order to
|
|
/// reconstruct the assistant's final text.
|
|
TextDelta(String),
|
|
/// Terminal event with the full assembled message (content + any
|
|
/// tool_calls). `message.content` equals the concatenation of every
|
|
/// preceding `TextDelta.0`.
|
|
Done {
|
|
message: ChatMessage,
|
|
prompt_eval_count: Option<i32>,
|
|
eval_count: Option<i32>,
|
|
},
|
|
}
|
|
|
|
/// Tool definition sent to the model (OpenAI-compatible function schema).
|
|
#[derive(Serialize, Clone, Debug)]
|
|
pub struct Tool {
|
|
#[serde(rename = "type")]
|
|
pub tool_type: String, // always "function"
|
|
pub function: ToolFunction,
|
|
}
|
|
|
|
#[derive(Serialize, Clone, Debug)]
|
|
pub struct ToolFunction {
|
|
pub name: String,
|
|
pub description: String,
|
|
pub parameters: serde_json::Value,
|
|
}
|
|
|
|
impl Tool {
|
|
pub fn function(name: &str, description: &str, parameters: serde_json::Value) -> Self {
|
|
Self {
|
|
tool_type: "function".to_string(),
|
|
function: ToolFunction {
|
|
name: name.to_string(),
|
|
description: description.to_string(),
|
|
parameters,
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A message in the chat conversation history.
|
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
|
pub struct ChatMessage {
|
|
pub role: String, // "system" | "user" | "assistant" | "tool"
|
|
/// Empty string (not null) when tool_calls is present — Ollama quirk.
|
|
#[serde(default)]
|
|
pub content: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub tool_calls: Option<Vec<ToolCall>>,
|
|
/// Base64 images — only on user messages to vision-capable models.
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub images: Option<Vec<String>>,
|
|
}
|
|
|
|
impl ChatMessage {
|
|
pub fn system(content: impl Into<String>) -> Self {
|
|
Self {
|
|
role: "system".to_string(),
|
|
content: content.into(),
|
|
tool_calls: None,
|
|
images: None,
|
|
}
|
|
}
|
|
pub fn user(content: impl Into<String>) -> Self {
|
|
Self {
|
|
role: "user".to_string(),
|
|
content: content.into(),
|
|
tool_calls: None,
|
|
images: None,
|
|
}
|
|
}
|
|
pub fn tool_result(content: impl Into<String>) -> Self {
|
|
Self {
|
|
role: "tool".to_string(),
|
|
content: content.into(),
|
|
tool_calls: None,
|
|
images: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Tool call returned by the model in an assistant message.
|
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
|
pub struct ToolCall {
|
|
pub function: ToolCallFunction,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub id: Option<String>,
|
|
}
|
|
|
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
|
pub struct ToolCallFunction {
|
|
pub name: String,
|
|
/// Canonical shape: native JSON. Providers that use JSON-encoded-string
|
|
/// arguments (OpenAI-compatible) translate at their wire boundary.
|
|
pub arguments: serde_json::Value,
|
|
}
|
|
|
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
|
pub struct ModelCapabilities {
|
|
pub name: String,
|
|
pub has_vision: bool,
|
|
pub has_tool_calling: bool,
|
|
/// Context window (tokens) the model is served with, when the server
|
|
/// fixes it regardless of the request's `num_ctx` (llama-swap) and it's
|
|
/// known without loading the model. Omitted otherwise.
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub context_window: Option<usize>,
|
|
}
|
|
|
|
/// A single node in the conversation tree. Wraps ChatMessage with structural
|
|
/// metadata so the tree can be persisted as JSON in `training_messages`.
|
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
|
pub struct StoredChatNode {
|
|
pub id: u64,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub parent_id: Option<u64>,
|
|
pub message: ChatMessage,
|
|
}
|
|
|
|
/// Top-level container for tree-based chat history. Stored as JSON in the
|
|
/// `training_messages` column. Backward-compatible reader detects whether
|
|
/// the stored value is a flat array (old format) or this object (new format).
|
|
#[derive(Serialize, Deserialize, Clone, Debug)]
|
|
pub struct ChatHistoryStore {
|
|
pub nodes: Vec<StoredChatNode>,
|
|
pub active_leaf_id: u64,
|
|
}
|
|
|
|
impl ChatHistoryStore {
|
|
/// Create an empty store. Test-only: production stores are always built
|
|
/// from `from_flat_array` or deserialized from `training_messages`.
|
|
#[cfg(test)]
|
|
pub fn empty() -> Self {
|
|
Self {
|
|
nodes: Vec::new(),
|
|
active_leaf_id: 0,
|
|
}
|
|
}
|
|
|
|
/// Convert a flat `Vec<ChatMessage>` (old format) into a tree.
|
|
/// Each message becomes a node chained by parent_id.
|
|
pub fn from_flat_array(messages: Vec<ChatMessage>) -> Self {
|
|
let mut nodes = Vec::with_capacity(messages.len());
|
|
let mut prev_id: Option<u64> = None;
|
|
for (i, msg) in messages.into_iter().enumerate() {
|
|
let id = (i + 1) as u64;
|
|
nodes.push(StoredChatNode {
|
|
id,
|
|
parent_id: prev_id,
|
|
message: msg,
|
|
});
|
|
prev_id = Some(id);
|
|
}
|
|
let active_leaf_id = prev_id.unwrap_or(0);
|
|
Self {
|
|
nodes,
|
|
active_leaf_id,
|
|
}
|
|
}
|
|
|
|
/// Allocate the next node ID (max existing + 1, or 1 if empty).
|
|
pub fn next_id(&self) -> u64 {
|
|
self.nodes.iter().map(|n| n.id).max().unwrap_or(0) + 1
|
|
}
|
|
|
|
/// Append a new node as a child of the given parent. Returns the new node's id.
|
|
pub fn append_node(&mut self, parent_id: Option<u64>, message: ChatMessage) -> u64 {
|
|
let id = self.next_id();
|
|
self.nodes.push(StoredChatNode {
|
|
id,
|
|
parent_id,
|
|
message,
|
|
});
|
|
id
|
|
}
|
|
|
|
/// Traverse from root to the given leaf_id, returning the path in order.
|
|
/// Returns None if the leaf_id doesn't exist.
|
|
pub fn path_to_leaf(&self, leaf_id: u64) -> Option<Vec<&StoredChatNode>> {
|
|
// Build a map from id to node for fast lookup.
|
|
let node_map: std::collections::HashMap<u64, usize> = self
|
|
.nodes
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, n)| (n.id, i))
|
|
.collect();
|
|
let &_leaf_idx = node_map.get(&leaf_id)?;
|
|
|
|
let mut path_indices = Vec::new();
|
|
let mut current_id = leaf_id;
|
|
while let Some(&idx) = node_map.get(¤t_id) {
|
|
path_indices.push(idx);
|
|
match self.nodes[idx].parent_id {
|
|
Some(pid) => current_id = pid,
|
|
None => break,
|
|
}
|
|
}
|
|
path_indices.reverse();
|
|
Some(path_indices.into_iter().map(|i| &self.nodes[i]).collect())
|
|
}
|
|
|
|
/// All direct children of a given node, ordered by id (chronological).
|
|
pub fn children_of(&self, node_id: u64) -> Vec<&StoredChatNode> {
|
|
self.nodes
|
|
.iter()
|
|
.filter(|n| n.parent_id == Some(node_id))
|
|
.collect()
|
|
}
|
|
|
|
/// For a given node, return (position, total) where:
|
|
/// - total = number of children of the parent of this node (i.e., siblings + self)
|
|
/// - position = 1-based rank of this node among siblings, ordered by id
|
|
///
|
|
/// Returns None if the node has no parent (root) or the parent has 0 or 1 child.
|
|
pub fn fork_at_node(&self, node_id: u64) -> Option<(usize, usize)> {
|
|
let node = self.nodes.iter().find(|n| n.id == node_id)?;
|
|
let parent_id = node.parent_id?;
|
|
let children = self.children_of(parent_id);
|
|
let total = children.len();
|
|
if total <= 1 {
|
|
return None;
|
|
}
|
|
let position = children.iter().position(|c| c.id == node_id)? + 1;
|
|
Some((position, total))
|
|
}
|
|
|
|
/// All leaf node IDs within the subtree rooted at `node_id` (the node
|
|
/// itself when it has no children). Empty when the node doesn't exist.
|
|
pub fn descendant_leaves(&self, node_id: u64) -> Vec<u64> {
|
|
if !self.nodes.iter().any(|n| n.id == node_id) {
|
|
return Vec::new();
|
|
}
|
|
let mut result = Vec::new();
|
|
let mut stack = vec![node_id];
|
|
while let Some(id) = stack.pop() {
|
|
let children = self.children_of(id);
|
|
if children.is_empty() {
|
|
result.push(id);
|
|
} else {
|
|
stack.extend(children.iter().map(|c| c.id));
|
|
}
|
|
}
|
|
result
|
|
}
|
|
|
|
/// Branch options at a divergence point: one entry per direct child of
|
|
/// `fork_node_id`, ordered by id (chronological — matches the stable
|
|
/// position rank shown in "X/Y" fork chips).
|
|
///
|
|
/// Each entry's `id` is the leaf to switch to for that branch: the
|
|
/// `preferred_leaf` when it lives in that child's subtree (so the
|
|
/// caller's current view maps onto its own option and can be filtered
|
|
/// out client-side), otherwise the most recent (max-id) descendant leaf.
|
|
///
|
|
/// Snippets preview the first visible message that *differs* between
|
|
/// the branches, not blindly the first child: Rewind & Regenerate
|
|
/// resends the identical question, so the first child is often shared
|
|
/// across siblings and would make every option read the same.
|
|
pub fn branch_options_at(&self, fork_node_id: u64, preferred_leaf: u64) -> Vec<BranchLeafInfo> {
|
|
struct Candidate {
|
|
rep: u64,
|
|
message_count: usize,
|
|
/// User/assistant contents after the divergence, tool
|
|
/// scaffolding (tool results, empty dispatch turns) skipped.
|
|
visible: Vec<String>,
|
|
}
|
|
let candidates: Vec<Candidate> = self
|
|
.children_of(fork_node_id)
|
|
.iter()
|
|
.map(|child| {
|
|
let leaves = self.descendant_leaves(child.id);
|
|
let rep = if leaves.contains(&preferred_leaf) {
|
|
preferred_leaf
|
|
} else {
|
|
leaves.iter().copied().max().unwrap_or(child.id)
|
|
};
|
|
let path = self.path_to_leaf(rep).unwrap_or_default();
|
|
let message_count = path.len();
|
|
let visible: Vec<String> = path
|
|
.iter()
|
|
.skip_while(|n| n.id != child.id)
|
|
.filter(|n| {
|
|
(n.message.role == "user" || n.message.role == "assistant")
|
|
&& !n.message.content.trim().is_empty()
|
|
})
|
|
.map(|n| n.message.content.clone())
|
|
.collect();
|
|
Candidate {
|
|
rep,
|
|
message_count,
|
|
visible,
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
// Walk forward past positions where every branch has the same
|
|
// content; stop at the first divergence (or when any branch runs
|
|
// out — a length difference is itself distinguishing).
|
|
let mut snippet_idx = 0usize;
|
|
loop {
|
|
let contents: Vec<Option<&String>> = candidates
|
|
.iter()
|
|
.map(|c| c.visible.get(snippet_idx))
|
|
.collect();
|
|
let all_present_and_equal =
|
|
contents.iter().all(|c| c.is_some()) && contents.windows(2).all(|w| w[0] == w[1]);
|
|
if !all_present_and_equal {
|
|
break;
|
|
}
|
|
snippet_idx += 1;
|
|
}
|
|
|
|
candidates
|
|
.into_iter()
|
|
.enumerate()
|
|
.map(|(i, c)| {
|
|
// A branch exhausted before the divergence index falls back
|
|
// to its own last message so it still shows *something*.
|
|
let snippet = c
|
|
.visible
|
|
.get(snippet_idx)
|
|
.or_else(|| c.visible.last())
|
|
.map(|s| s.chars().take(100).collect::<String>())
|
|
.unwrap_or_default();
|
|
BranchLeafInfo {
|
|
id: c.rep,
|
|
snippet,
|
|
message_count: c.message_count,
|
|
position: Some(i + 1),
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// All leaf node IDs (nodes with no children).
|
|
pub fn leaves(&self) -> Vec<u64> {
|
|
let child_ids: std::collections::HashSet<u64> =
|
|
self.nodes.iter().filter_map(|n| n.parent_id).collect();
|
|
self.nodes
|
|
.iter()
|
|
.filter(|n| !child_ids.contains(&n.id))
|
|
.map(|n| n.id)
|
|
.collect()
|
|
}
|
|
|
|
/// The parent of the given node, if any. Used to re-anchor a rewind
|
|
/// that discards every rendered message onto the seed node above them.
|
|
pub fn parent_of(&self, node_id: u64) -> Option<&StoredChatNode> {
|
|
let node = self.nodes.iter().find(|n| n.id == node_id)?;
|
|
let parent_id = node.parent_id?;
|
|
self.nodes.iter().find(|n| n.id == parent_id)
|
|
}
|
|
|
|
/// Return all leaf nodes with metadata: id, snippet (first 100 chars of
|
|
/// the first user-visible message after the divergence point), and
|
|
/// message_count (number of nodes in the path from root to leaf).
|
|
pub fn leaves_with_info(&self) -> Vec<BranchLeafInfo> {
|
|
let leaf_ids = self.leaves();
|
|
leaf_ids
|
|
.into_iter()
|
|
.map(|leaf_id| {
|
|
let path = self.path_to_leaf(leaf_id).unwrap_or_default();
|
|
let message_count = path.len();
|
|
// Snippet: first 100 chars of the first user/assistant message
|
|
// (skip system messages).
|
|
let snippet = path
|
|
.iter()
|
|
.find(|n| n.message.role != "system")
|
|
.map(|n| n.message.content.chars().take(100).collect::<String>())
|
|
.unwrap_or_default();
|
|
BranchLeafInfo {
|
|
id: leaf_id,
|
|
snippet,
|
|
message_count,
|
|
position: None,
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
/// Metadata about a branch leaf, returned by the branches endpoint.
|
|
#[derive(Serialize, Clone, Debug)]
|
|
pub struct BranchLeafInfo {
|
|
pub id: u64,
|
|
#[serde(skip_serializing_if = "String::is_empty")]
|
|
pub snippet: String,
|
|
pub message_count: usize,
|
|
/// 1-based sibling rank at the divergence point. Only present for
|
|
/// scoped queries (`branch_options_at`); tree-wide leaf listings have
|
|
/// no single divergence to rank against.
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub position: Option<usize>,
|
|
}
|
|
|
|
/// Strip a leading `<think>…</think>` reasoning block from model output.
|
|
///
|
|
/// Thinking models sometimes emit chain-of-thought inside think tags before
|
|
/// the real answer. Everything after the first `</think>` is the answer;
|
|
/// when no tag is present — or the text after it is empty — the trimmed
|
|
/// input is returned unchanged. Mirrors the behavior Ollama's
|
|
/// `extract_final_answer` has applied to single-shot generation; shared here
|
|
/// so the tool-calling final-content paths (agentic generation + chat) can
|
|
/// apply the identical cleanup before parsing / persisting.
|
|
pub fn strip_think_blocks(response: &str) -> String {
|
|
let response = response.trim();
|
|
|
|
if let Some(pos) = response.find("</think>") {
|
|
let answer = response[pos + "</think>".len()..].trim();
|
|
if !answer.is_empty() {
|
|
return answer.to_string();
|
|
}
|
|
}
|
|
|
|
response.to_string()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn strip_think_blocks_removes_leading_think_block() {
|
|
let raw = "<think>\nLet me reason about this.\n</think>\n\nTitle: A Day Out\n\nThe body.";
|
|
assert_eq!(strip_think_blocks(raw), "Title: A Day Out\n\nThe body.");
|
|
}
|
|
|
|
#[test]
|
|
fn strip_think_blocks_passes_through_plain_content() {
|
|
assert_eq!(strip_think_blocks(" just an answer "), "just an answer");
|
|
}
|
|
|
|
#[test]
|
|
fn strip_think_blocks_keeps_content_when_answer_after_tag_is_empty() {
|
|
// A think block with nothing after it: better to return the trimmed
|
|
// original than an empty string (matches Ollama's fallback).
|
|
let raw = "<think>only thoughts</think>";
|
|
assert_eq!(strip_think_blocks(raw), raw);
|
|
}
|
|
|
|
#[test]
|
|
fn strip_think_blocks_handles_unclosed_tag() {
|
|
let raw = "<think>thinking forever";
|
|
assert_eq!(strip_think_blocks(raw), raw);
|
|
}
|
|
|
|
// ── ChatHistoryStore tree traversal tests ──────────────────────────
|
|
|
|
fn mk_msg(role: &str, content: &str) -> ChatMessage {
|
|
match role {
|
|
"system" => ChatMessage::system(content),
|
|
"user" => ChatMessage::user(content),
|
|
_ => ChatMessage {
|
|
role: role.to_string(),
|
|
content: content.to_string(),
|
|
tool_calls: None,
|
|
images: None,
|
|
},
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn store_from_flat_array_chains_parent_ids() {
|
|
let msgs = vec![
|
|
mk_msg("system", "sys"),
|
|
mk_msg("user", "hello"),
|
|
mk_msg("assistant", "hi there"),
|
|
];
|
|
let store = ChatHistoryStore::from_flat_array(msgs);
|
|
assert_eq!(store.nodes.len(), 3);
|
|
assert_eq!(store.nodes[0].id, 1);
|
|
assert_eq!(store.nodes[0].parent_id, None);
|
|
assert_eq!(store.nodes[1].parent_id, Some(1));
|
|
assert_eq!(store.nodes[2].parent_id, Some(2));
|
|
assert_eq!(store.active_leaf_id, 3);
|
|
}
|
|
|
|
#[test]
|
|
fn store_path_to_leaf_returns_ordered_path() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let n1 = store.append_node(None, mk_msg("user", "u1"));
|
|
let n2 = store.append_node(Some(n1), mk_msg("assistant", "a1"));
|
|
let n3 = store.append_node(Some(n2), mk_msg("user", "u2"));
|
|
let path = store.path_to_leaf(n3).unwrap();
|
|
assert_eq!(path.len(), 3);
|
|
assert_eq!(path[0].id, n1);
|
|
assert_eq!(path[1].id, n2);
|
|
assert_eq!(path[2].id, n3);
|
|
}
|
|
|
|
#[test]
|
|
fn store_path_to_leaf_returns_none_for_missing_id() {
|
|
let store = ChatHistoryStore::empty();
|
|
assert!(store.path_to_leaf(999).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn store_children_of_returns_direct_children() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("user", "root"));
|
|
let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
|
|
let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
|
|
let _d1 = store.append_node(Some(c1), mk_msg("assistant", "d1"));
|
|
let children = store.children_of(root);
|
|
assert_eq!(children.len(), 2);
|
|
assert_eq!(children[0].id, c1);
|
|
assert_eq!(children[1].id, c2);
|
|
}
|
|
|
|
#[test]
|
|
fn store_fork_at_node_detects_fork() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("user", "root"));
|
|
let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
|
|
let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
|
|
let c3 = store.append_node(Some(root), mk_msg("user", "c3"));
|
|
// c1 is position 1 of 3 siblings
|
|
assert_eq!(store.fork_at_node(c1), Some((1, 3)));
|
|
assert_eq!(store.fork_at_node(c2), Some((2, 3)));
|
|
assert_eq!(store.fork_at_node(c3), Some((3, 3)));
|
|
// root has no parent, so no fork
|
|
assert!(store.fork_at_node(root).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn store_fork_at_node_returns_none_for_single_child() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("user", "root"));
|
|
let _only = store.append_node(Some(root), mk_msg("assistant", "only"));
|
|
assert!(store.fork_at_node(2).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn store_leaves_returns_leaf_ids() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("user", "root"));
|
|
let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
|
|
let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
|
|
let _d1 = store.append_node(Some(c1), mk_msg("assistant", "d1"));
|
|
let leaves = store.leaves();
|
|
assert_eq!(leaves.len(), 2);
|
|
assert!(leaves.contains(&c2));
|
|
assert!(leaves.contains(&_d1));
|
|
assert!(!leaves.contains(&root));
|
|
assert!(!leaves.contains(&c1));
|
|
}
|
|
|
|
#[test]
|
|
fn store_parent_of_returns_parent() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("user", "root"));
|
|
let child = store.append_node(Some(root), mk_msg("assistant", "child"));
|
|
let parent = store.parent_of(child).unwrap();
|
|
assert_eq!(parent.id, root);
|
|
assert!(store.parent_of(root).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn store_next_id_increments() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
assert_eq!(store.next_id(), 1);
|
|
store.append_node(None, mk_msg("user", "a"));
|
|
assert_eq!(store.next_id(), 2);
|
|
store.append_node(Some(1), mk_msg("assistant", "b"));
|
|
assert_eq!(store.next_id(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn store_descendant_leaves_walks_subtree() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("user", "root"));
|
|
let c1 = store.append_node(Some(root), mk_msg("assistant", "c1"));
|
|
let c2 = store.append_node(Some(root), mk_msg("user", "c2"));
|
|
let d1 = store.append_node(Some(c1), mk_msg("assistant", "d1"));
|
|
let d2 = store.append_node(Some(c1), mk_msg("assistant", "d2"));
|
|
// c1's subtree has two leaves; c2 is its own leaf; root sees all.
|
|
let mut c1_leaves = store.descendant_leaves(c1);
|
|
c1_leaves.sort_unstable();
|
|
assert_eq!(c1_leaves, vec![d1, d2]);
|
|
assert_eq!(store.descendant_leaves(c2), vec![c2]);
|
|
let mut all = store.descendant_leaves(root);
|
|
all.sort_unstable();
|
|
assert_eq!(all, vec![c2, d1, d2]);
|
|
assert!(store.descendant_leaves(999).is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn branch_options_ranks_children_and_prefers_viewing_leaf() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("user", "root"));
|
|
// First branch: two leaves (older + newer).
|
|
let c1 = store.append_node(Some(root), mk_msg("assistant", "reply one"));
|
|
let old_leaf = store.append_node(Some(c1), mk_msg("user", "old follow-up"));
|
|
let new_leaf = store.append_node(Some(c1), mk_msg("user", "new follow-up"));
|
|
// Second branch: single leaf.
|
|
let c2 = store.append_node(Some(root), mk_msg("assistant", "reply two"));
|
|
|
|
// Viewing the *older* leaf of branch 1: its option must anchor to
|
|
// that exact leaf, not the max-id one, so the client's
|
|
// "hide the branch I'm on" filter matches.
|
|
let opts = store.branch_options_at(root, old_leaf);
|
|
assert_eq!(opts.len(), 2);
|
|
assert_eq!(opts[0].id, old_leaf);
|
|
assert_eq!(opts[0].position, Some(1));
|
|
assert_eq!(opts[0].snippet, "reply one");
|
|
assert_eq!(opts[1].id, c2);
|
|
assert_eq!(opts[1].position, Some(2));
|
|
|
|
// Viewing branch 2: branch 1 falls back to its most recent leaf.
|
|
let opts = store.branch_options_at(root, c2);
|
|
assert_eq!(opts[0].id, new_leaf);
|
|
assert_eq!(opts[1].id, c2);
|
|
}
|
|
|
|
#[test]
|
|
fn branch_options_snippet_skips_shared_resent_question() {
|
|
// Rewind & Regenerate resends the identical question, so both
|
|
// branches' first child matches — the preview must advance to the
|
|
// first *differing* message (the assistant replies) or every
|
|
// option would read the same.
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("assistant", "prior reply"));
|
|
let q1 = store.append_node(Some(root), mk_msg("user", "same question"));
|
|
let _a1 = store.append_node(Some(q1), mk_msg("assistant", "first answer"));
|
|
let q2 = store.append_node(Some(root), mk_msg("user", "same question"));
|
|
let a2 = store.append_node(Some(q2), mk_msg("assistant", "second answer"));
|
|
|
|
let opts = store.branch_options_at(root, a2);
|
|
assert_eq!(opts.len(), 2);
|
|
assert_eq!(opts[0].snippet, "first answer");
|
|
assert_eq!(opts[1].snippet, "second answer");
|
|
}
|
|
|
|
#[test]
|
|
fn branch_options_snippet_falls_back_when_branch_exhausted() {
|
|
// One branch is just the shared question (no reply yet); the other
|
|
// continues past it. The shorter branch falls back to its own last
|
|
// message rather than showing an empty preview.
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("assistant", "prior reply"));
|
|
let _q1 = store.append_node(Some(root), mk_msg("user", "same question"));
|
|
let q2 = store.append_node(Some(root), mk_msg("user", "same question"));
|
|
let a2 = store.append_node(Some(q2), mk_msg("assistant", "second answer"));
|
|
|
|
let opts = store.branch_options_at(root, a2);
|
|
assert_eq!(opts[0].snippet, "same question");
|
|
assert_eq!(opts[1].snippet, "second answer");
|
|
}
|
|
|
|
#[test]
|
|
fn branch_options_snippet_skips_tool_scaffolding() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("user", "root"));
|
|
// Branch whose diverging child is an empty tool-dispatch turn.
|
|
let dispatch = store.append_node(
|
|
Some(root),
|
|
ChatMessage {
|
|
role: "assistant".to_string(),
|
|
content: String::new(),
|
|
tool_calls: None,
|
|
images: None,
|
|
},
|
|
);
|
|
let tool = store.append_node(Some(dispatch), mk_msg("tool", "raw tool output"));
|
|
let reply = store.append_node(Some(tool), mk_msg("assistant", "final answer"));
|
|
let _other = store.append_node(Some(root), mk_msg("assistant", "direct reply"));
|
|
|
|
let opts = store.branch_options_at(root, reply);
|
|
assert_eq!(opts[0].snippet, "final answer");
|
|
assert_eq!(opts[1].snippet, "direct reply");
|
|
}
|
|
|
|
#[test]
|
|
fn store_append_node_sets_parent() {
|
|
let mut store = ChatHistoryStore::empty();
|
|
let parent = store.append_node(None, mk_msg("user", "parent"));
|
|
let child = store.append_node(Some(parent), mk_msg("assistant", "child"));
|
|
assert_eq!(store.nodes.len(), 2);
|
|
assert_eq!(store.nodes[1].parent_id, Some(parent));
|
|
assert_eq!(store.nodes[1].id, child);
|
|
}
|
|
|
|
#[test]
|
|
fn branch_options_handles_nested_forks() {
|
|
// Regression: verify sub-fork display when there are multiple forks.
|
|
// Structure:
|
|
// root (user "q1")
|
|
// ├── a1 (assistant "answer 1")
|
|
// └── a2 (assistant "answer 2")
|
|
// ├── q2 (user "q2")
|
|
// │ ├── r1 (assistant "reply 1")
|
|
// │ └── r2 (assistant "reply 2") <- active leaf
|
|
//
|
|
// branch_options_at(root) should return [a1, a2] (2 branches)
|
|
// branch_options_at(q2) should return [r1, r2] (2 sub-branches)
|
|
let mut store = ChatHistoryStore::empty();
|
|
let root = store.append_node(None, mk_msg("user", "q1"));
|
|
let a1 = store.append_node(Some(root), mk_msg("assistant", "answer 1"));
|
|
let a2 = store.append_node(Some(root), mk_msg("assistant", "answer 2"));
|
|
let q2 = store.append_node(Some(a2), mk_msg("user", "q2"));
|
|
let r1 = store.append_node(Some(q2), mk_msg("assistant", "reply 1"));
|
|
let r2 = store.append_node(Some(q2), mk_msg("assistant", "reply 2"));
|
|
store.active_leaf_id = r2;
|
|
|
|
// Root fork: should show [a1, a2's subtree]
|
|
let opts = store.branch_options_at(root, r2);
|
|
assert_eq!(opts.len(), 2);
|
|
assert_eq!(opts[0].id, a1);
|
|
assert_eq!(opts[0].position, Some(1));
|
|
assert_eq!(opts[0].snippet, "answer 1");
|
|
// a2's subtree includes r2, so the rep should be r2 (preferred leaf)
|
|
assert_eq!(opts[1].id, r2);
|
|
assert_eq!(opts[1].position, Some(2));
|
|
assert_eq!(opts[1].snippet, "answer 2");
|
|
|
|
// Sub-fork at q2: should show [r1, r2]
|
|
let opts = store.branch_options_at(q2, r2);
|
|
assert_eq!(opts.len(), 2);
|
|
assert_eq!(opts[0].id, r1);
|
|
assert_eq!(opts[0].position, Some(1));
|
|
assert_eq!(opts[0].snippet, "reply 1");
|
|
assert_eq!(opts[1].id, r2);
|
|
assert_eq!(opts[1].position, Some(2));
|
|
assert_eq!(opts[1].snippet, "reply 2");
|
|
}
|
|
}
|