use std::sync::Arc;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;
use super::Middleware;
use super::MiddlewareCommandContext;
use super::MiddlewareCommandOutput;
use super::RuntimeContext;
use super::attachments::strip_attachment_references;
use super::manifest::{MiddlewareManifest, MiddlewareSettingManifest};
use super::tools::{
Catalog, ExecutionMode, Tool, ToolContext, ToolExposure, rank_bm25, render_tool_event,
};
use crate::BoxFuture;
use crate::Error;
use crate::Result;
use crate::backend::checkpoint::Checkpoint;
use crate::backend::checkpoint::SessionCursor;
use crate::backend::checkpoint::SessionPage;
use crate::backend::checkpoint::SessionPageRequest;
use crate::backend::checkpoint::SessionSummary;
use crate::backend::checkpoint::TranscriptPageRequest;
use crate::backend::model::ToolDefinition;
use crate::protocol::EventMsg;
use crate::protocol::FrontendBlock;
use crate::protocol::FrontendCommand;
use crate::protocol::FrontendContribution;
use crate::protocol::FrontendEvent;
use crate::protocol::FrontendPickerOption;
use crate::protocol::FrontendSlot;
use crate::protocol::FrontendSymbol;
use crate::protocol::FrontendTone;
use crate::protocol::FrontendWidget;
use crate::protocol::MessageTarget;
use crate::protocol::Op;
use crate::protocol::replay_events;
mod text {
include!(concat!(env!("OUT_DIR"), "/src_middleware_sessions_text.rs"));
}
const MAX_PAGE_SIZE: usize = 1_000;
const MAX_HISTORY_QUERY_BYTES: usize = 512;
const MAX_HISTORY_CURSOR_BYTES: usize = 8_192;
const MAX_HISTORY_RESULTS: usize = 6;
const MAX_HISTORY_PAGES: usize = 32;
const MAX_HISTORY_ITEMS: usize = 128;
const MAX_HISTORY_SCAN_CHARS: usize = 64_000;
const HISTORY_CHUNK_CHARS: usize = 8_000;
const HISTORY_EXCERPT_CHARS: usize = 600;
const MAX_HISTORY_READ_CHARS: usize = 4_000;
const _: () = {
assert!(text::DEFAULTS_PAGE_SIZE >= 1);
assert!(text::DEFAULTS_PAGE_SIZE <= MAX_PAGE_SIZE as i64);
assert!(text::SETTING_PAGE_SIZE_STEP > 0);
};
pub const DEFAULT_PAGE_SIZE: usize = text::DEFAULTS_PAGE_SIZE as usize;
const SETTINGS: &[MiddlewareSettingManifest] = &[MiddlewareSettingManifest::Integer {
id: "page_size",
label: text::SETTING_PAGE_SIZE_LABEL,
description: text::SETTING_PAGE_SIZE_DESCRIPTION,
min: 1,
max: Some(MAX_PAGE_SIZE as i64),
step: text::SETTING_PAGE_SIZE_STEP,
default: DEFAULT_PAGE_SIZE as i64,
}];
pub const MANIFEST: MiddlewareManifest = MiddlewareManifest {
id: "sessions",
label: text::MANIFEST_LABEL,
description: text::MANIFEST_DESCRIPTION,
required: true,
default_enabled: true,
settings: SETTINGS,
};
pub struct Sessions {
page_size: usize,
}
impl Sessions {
pub fn new(page_size: usize) -> Result<Self> {
if page_size == 0 || page_size > MAX_PAGE_SIZE {
return Err(Error::Config(format!(
"chat catalog page size must be between 1 and {MAX_PAGE_SIZE}"
)));
}
Ok(Self { page_size })
}
}
impl Default for Sessions {
fn default() -> Self {
Self {
page_size: DEFAULT_PAGE_SIZE,
}
}
}
impl Middleware for Sessions {
fn name(&self) -> &'static str {
MANIFEST.id
}
fn register(&self, catalog: &mut Catalog, runtime: &RuntimeContext) -> Result<()> {
let history = Arc::new(History {
checkpoints: Arc::clone(&runtime.checkpoints),
session_id: runtime.session_id.clone(),
bot_id: runtime.session_context.bot_id.clone(),
});
catalog.register(Arc::new(SearchHistory(Arc::clone(&history))))?;
catalog.register(Arc::new(ReadHistory(history)))
}
fn frontend(&self) -> FrontendContribution {
FrontendContribution {
capability: self.name().into(),
accepts_file_attachments: false,
count: None,
commands: vec![
FrontendCommand {
name: "resume".into(),
arguments: String::new(),
description: text::COMMAND_RESUME_DESCRIPTION.into(),
requires_idle: true,
},
FrontendCommand {
name: "fork".into(),
arguments: String::new(),
description: text::COMMAND_FORK_DESCRIPTION.into(),
requires_idle: true,
},
],
widgets: vec![FrontendWidget {
id: "fork".into(),
slot: FrontendSlot::MessageActions,
text: text::WIDGET_FORK_CHAT.into(),
tone: FrontendTone::Neutral,
symbol: Some(FrontendSymbol::Branch),
icon_only: true,
progress: None,
content: None,
action: Some(Op::CapabilityCommand {
capability: MANIFEST.id.into(),
command: "fork".into(),
arguments: String::new(),
input: None,
target: None,
}),
}],
references: Vec::new(),
}
}
fn render(&self, event: &EventMsg, _session_id: &str) -> Option<FrontendBlock> {
render_tool_event(
event,
|name| matches!(name, "search_history" | "read_history"),
|name, arguments| super::tools::ToolHeading {
title: if matches!(event, EventMsg::ToolCallEnd(_)) {
name
} else {
match name {
"search_history" => text::RENDER_SEARCH_HISTORY,
_ => text::RENDER_READ_HISTORY,
}
}
.into(),
detail: arguments
.get("query")
.and_then(Value::as_str)
.unwrap_or_default()
.into(),
},
)
}
fn command<'a>(
&'a self,
context: MiddlewareCommandContext<'a>,
) -> BoxFuture<'a, Result<MiddlewareCommandOutput>> {
Box::pin(async move {
match context.command {
"resume" => resume(context, self.page_size).await,
"fork" => fork(context).await,
command => Err(Error::Unknown(format!("sessions command `{command}`"))),
}
})
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
enum HistoryScope {
#[default]
Current,
OtherChats,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct SearchHistoryArgs {
query: String,
#[serde(default)]
scope: HistoryScope,
cursor: Option<String>,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ReadHistoryArgs {
session_id: Option<String>,
target: MessageTarget,
#[serde(default)]
offset: usize,
#[serde(default = "default_read_chars")]
max_chars: usize,
}
const fn default_read_chars() -> usize {
4_000
}
#[derive(Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
struct HistoryCursor {
scope: HistoryScope,
query: String,
catalog: Option<SessionCursor>,
session_id: Option<String>,
before_sequence: Option<u64>,
remaining_items: Option<usize>,
offset: usize,
}
#[derive(Serialize)]
struct HistoryHit {
session_id: String,
target: MessageTarget,
kind: &'static str,
offset: usize,
excerpt: String,
}
struct HistoryDocument {
session_id: String,
target: MessageTarget,
kind: &'static str,
offset: usize,
text: String,
}
struct History {
checkpoints: Arc<dyn crate::backend::checkpoint::CheckpointStore>,
session_id: String,
bot_id: String,
}
struct SearchHistory(Arc<History>);
struct ReadHistory(Arc<History>);
impl Tool for SearchHistory {
fn definition(&self) -> ToolDefinition {
ToolDefinition {
name: "search_history".into(),
description: text::TOOL_SEARCH_HISTORY_DESCRIPTION.into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"query": {"type": "string", "maxLength": MAX_HISTORY_QUERY_BYTES,
"description": text::TOOL_SEARCH_HISTORY_PARAMETER_QUERY_DESCRIPTION},
"scope": {"type": "string", "enum": ["current", "other_chats"],
"description": "Defaults to this chat. other_chats explicitly searches this Bot's other chats."},
"cursor": {"type": "string", "maxLength": MAX_HISTORY_CURSOR_BYTES,
"description": "Unmodified next_cursor from the same query and scope; continue even when hits is empty."}
},
"required": ["query"],
"additionalProperties": false
}),
}
}
fn exposure(&self) -> ToolExposure {
ToolExposure::Direct
}
fn execution_mode(&self) -> ExecutionMode {
ExecutionMode::Parallel
}
fn call<'a>(
&'a self,
_context: ToolContext,
arguments: Value,
) -> BoxFuture<'a, Result<String>> {
Box::pin(async move { self.0.search(serde_json::from_value(arguments)?).await })
}
}
impl Tool for ReadHistory {
fn definition(&self) -> ToolDefinition {
ToolDefinition {
name: "read_history".into(),
description: text::TOOL_READ_HISTORY_DESCRIPTION.into(),
parameters: serde_json::json!({
"type": "object",
"properties": {
"session_id": {"type": "string", "maxLength": 512,
"description": "Defaults to this chat; an explicit other chat must belong to this Bot."},
"target": {"type": "object", "properties": {
"checkpoint_sequence": {"type": "integer", "minimum": 0},
"batch_item_count": {"type": "integer", "minimum": 1}
}, "required": ["checkpoint_sequence", "batch_item_count"], "additionalProperties": false},
"offset": {"type": "integer", "minimum": 0,
"description": "Character offset, initially zero or the search hit's offset."},
"max_chars": {"type": "integer", "minimum": 1, "maximum": MAX_HISTORY_READ_CHARS,
"description": "Maximum characters to return; defaults to 4000. Continue at next_offset."}
},
"required": ["target"],
"additionalProperties": false
}),
}
}
fn exposure(&self) -> ToolExposure {
ToolExposure::Direct
}
fn execution_mode(&self) -> ExecutionMode {
ExecutionMode::Parallel
}
fn call<'a>(
&'a self,
_context: ToolContext,
arguments: Value,
) -> BoxFuture<'a, Result<String>> {
Box::pin(async move { self.0.read(serde_json::from_value(arguments)?).await })
}
}
impl History {
async fn authorize(&self, session_id: &str) -> Result<()> {
validate_history_session_id(session_id)?;
let checkpoint = self.checkpoints.load(session_id).await?;
if !checkpoint.is_some_and(|checkpoint| checkpoint.session_context.bot_id == self.bot_id) {
return Err(Error::Tool(
"history chat is unavailable to this Bot".into(),
));
}
Ok(())
}
fn cursor(&self, arguments: SearchHistoryArgs) -> Result<HistoryCursor> {
let query = arguments.query.trim();
if query.is_empty() || query.len() > MAX_HISTORY_QUERY_BYTES {
return Err(Error::Tool(format!(
"history query must be 1–{MAX_HISTORY_QUERY_BYTES} bytes"
)));
}
let Some(value) = arguments.cursor else {
return Ok(HistoryCursor {
scope: arguments.scope,
query: query.into(),
catalog: None,
session_id: (arguments.scope == HistoryScope::Current)
.then(|| self.session_id.clone()),
before_sequence: None,
remaining_items: None,
offset: 0,
});
};
if value.len() > MAX_HISTORY_CURSOR_BYTES {
return Err(Error::Tool("history cursor is too large".into()));
}
let cursor: HistoryCursor = serde_json::from_str(&value)?;
if cursor.scope != arguments.scope
|| cursor.query != query
|| cursor.remaining_items == Some(0)
|| (cursor.remaining_items.is_some() && cursor.before_sequence.is_none())
|| (cursor.remaining_items.is_none() && cursor.offset != 0)
|| (cursor.scope == HistoryScope::Current
&& (cursor.session_id.as_deref() != Some(&self.session_id)
|| cursor.catalog.is_some()))
|| (cursor.scope == HistoryScope::OtherChats
&& cursor.session_id.as_deref() == Some(&self.session_id))
{
return Err(Error::Tool(
"history cursor does not match this query and scope".into(),
));
}
if let Some(catalog) = &cursor.catalog {
validate_history_session_id(&catalog.session_id)?;
validate_history_sequence(catalog.sequence)?;
}
if let Some(sequence) = cursor.before_sequence {
validate_history_sequence(sequence)?;
}
Ok(cursor)
}
async fn search(&self, arguments: SearchHistoryArgs) -> Result<String> {
let mut cursor = self.cursor(arguments)?;
if let Some(session_id) = &cursor.session_id {
self.authorize(session_id).await?;
}
if let Some(catalog) = &cursor.catalog {
self.authorize(&catalog.session_id).await?;
}
let mut documents = Vec::new();
let mut scanned_chars = 0;
let mut scanned_items = 0;
let mut complete = false;
for _ in 0..MAX_HISTORY_PAGES {
let Some(session_id) = cursor.session_id.clone() else {
if !self.advance_session(&mut cursor).await? {
complete = true;
break;
}
continue;
};
let page = self
.checkpoints
.transcript_page(
&session_id,
TranscriptPageRequest {
before_sequence: cursor.before_sequence,
max_batches: 1,
},
)
.await?;
let Some(batch) = page.batches.into_iter().next() else {
if cursor.scope == HistoryScope::Current {
complete = true;
break;
}
cursor.session_id = None;
cursor.before_sequence = None;
cursor.remaining_items = None;
cursor.offset = 0;
continue;
};
scan_history_batch(
&batch,
&mut cursor,
&mut documents,
&mut scanned_items,
&mut scanned_chars,
)?;
if scanned_items >= MAX_HISTORY_ITEMS
|| MAX_HISTORY_SCAN_CHARS - scanned_chars <= MAX_HISTORY_QUERY_BYTES
{
break;
}
}
let searchable = documents
.iter()
.map(|document| document.text.clone())
.collect::<Vec<_>>();
let hits = rank_bm25(&searchable, &cursor.query, MAX_HISTORY_RESULTS)
.into_iter()
.map(|index| {
let document = &documents[index];
let (offset, excerpt) = history_excerpt(&document.text, &cursor.query);
HistoryHit {
session_id: document.session_id.clone(),
target: document.target,
kind: document.kind,
offset: document.offset + offset,
excerpt,
}
})
.collect::<Vec<_>>();
let next_cursor = (!complete)
.then(|| serde_json::to_string(&cursor))
.transpose()?;
Ok(serde_json::to_string(&serde_json::json!({
"hits": hits, "next_cursor": next_cursor,
"scanned_items": scanned_items, "scanned_chars": scanned_chars,
}))?)
}
async fn advance_session(&self, cursor: &mut HistoryCursor) -> Result<bool> {
let page = self
.checkpoints
.list_sessions_page(SessionPageRequest {
bot_id: Some(self.bot_id.clone()),
cursor: cursor.catalog.clone(),
limit: 1,
})
.await?;
let Some(summary) = page.sessions.into_iter().next() else {
return Ok(false);
};
if summary.session_context.bot_id != self.bot_id {
return Err(Error::Checkpoint(
"session catalog returned another Bot's chat".into(),
));
}
cursor.catalog = Some(SessionCursor {
updated_at: summary.updated_at,
sequence: summary.sequence,
session_id: summary.session_id.clone(),
});
if summary.session_id != self.session_id {
self.authorize(&summary.session_id).await?;
cursor.session_id = Some(summary.session_id);
}
Ok(true)
}
async fn read(&self, arguments: ReadHistoryArgs) -> Result<String> {
if arguments.max_chars == 0 || arguments.max_chars > MAX_HISTORY_READ_CHARS {
return Err(Error::Tool(format!(
"max_chars must be 1–{MAX_HISTORY_READ_CHARS}"
)));
}
let session_id = arguments.session_id.as_deref().unwrap_or(&self.session_id);
self.authorize(session_id).await?;
let before_sequence = arguments
.target
.checkpoint_sequence
.checked_add(1)
.ok_or_else(|| Error::Tool("history target sequence is too large".into()))?;
validate_history_sequence(before_sequence)?;
let page = self
.checkpoints
.transcript_page(
session_id,
TranscriptPageRequest {
before_sequence: Some(before_sequence),
max_batches: 1,
},
)
.await?;
let item = page
.batches
.first()
.filter(|batch| batch.sequence == arguments.target.checkpoint_sequence)
.and_then(|batch| {
arguments
.target
.batch_item_count
.checked_sub(1)
.and_then(|index| batch.items.get(index))
})
.and_then(history_text)
.ok_or_else(|| {
Error::Tool("history target is not a readable transcript item".into())
})?;
let (content, next_offset) = history_chunk(&item.1, arguments.offset, arguments.max_chars)?;
Ok(serde_json::to_string(&serde_json::json!({
"session_id": session_id, "target": arguments.target, "kind": item.0,
"offset": arguments.offset, "text": content, "next_offset": next_offset,
}))?)
}
}
fn validate_history_session_id(session_id: &str) -> Result<()> {
if session_id.trim().is_empty()
|| session_id.len() > 512
|| session_id.chars().any(char::is_control)
{
return Err(Error::Tool("history session_id must be 1–512 bytes".into()));
}
Ok(())
}
fn validate_history_sequence(sequence: u64) -> Result<()> {
if i64::try_from(sequence).is_err() {
return Err(Error::Tool(
"history sequence exceeds the supported range".into(),
));
}
Ok(())
}
fn scan_history_batch(
batch: &crate::backend::checkpoint::TranscriptBatch,
cursor: &mut HistoryCursor,
documents: &mut Vec<HistoryDocument>,
scanned_items: &mut usize,
scanned_chars: &mut usize,
) -> Result<()> {
let mut remaining = cursor.remaining_items.unwrap_or(batch.items.len());
if remaining > batch.items.len()
|| (cursor.remaining_items.is_some()
&& cursor.before_sequence != batch.sequence.checked_add(1))
{
return Err(Error::Tool(
"history cursor no longer identifies a transcript item".into(),
));
}
while remaining > 0
&& *scanned_items < MAX_HISTORY_ITEMS
&& MAX_HISTORY_SCAN_CHARS - *scanned_chars > MAX_HISTORY_QUERY_BYTES
{
*scanned_items += 1;
let mut next_offset = None;
if let Some((kind, text)) = history_text(&batch.items[remaining - 1]) {
let limit = HISTORY_CHUNK_CHARS.min(MAX_HISTORY_SCAN_CHARS - *scanned_chars);
let (chunk, next) = history_chunk(&text, cursor.offset, limit)?;
*scanned_chars += chunk.chars().count();
documents.push(HistoryDocument {
session_id: cursor
.session_id
.clone()
.ok_or_else(|| Error::Tool("history cursor has no chat".into()))?,
target: MessageTarget {
checkpoint_sequence: batch.sequence,
batch_item_count: remaining,
},
kind,
offset: cursor.offset,
text: chunk,
});
next_offset = next.map(|offset| offset - cursor.query.chars().count());
}
if let Some(offset) = next_offset {
cursor.offset = offset;
} else {
remaining -= 1;
cursor.offset = 0;
}
}
cursor.before_sequence = Some(if remaining == 0 {
batch.sequence
} else {
batch
.sequence
.checked_add(1)
.ok_or_else(|| Error::Tool("history sequence is too large".into()))?
});
cursor.remaining_items = (remaining > 0).then_some(remaining);
Ok(())
}
fn history_text(item: &Value) -> Option<(&'static str, String)> {
if let Some(message) = crate::protocol::message_metadata(item) {
return Some(("user", message.text));
}
if crate::protocol::is_internal_message(item) {
return None;
}
match item.get("type").and_then(Value::as_str) {
Some("function_call") => Some((
"tool_call",
format!(
"{}\n{}",
item.get("name")?.as_str()?,
history_value_text(item.get("arguments")?)
),
)),
Some("function_call_output") => {
Some(("tool_result", history_value_text(item.get("output")?)))
}
Some("reasoning" | "compaction") => None,
_ => {
let kind = match item.get("role")?.as_str()? {
"user" => "user",
"assistant" => "assistant",
_ => return None,
};
let content = item.get("content")?;
let text = match content {
Value::String(text) => text.clone(),
Value::Array(parts) => parts
.iter()
.filter_map(|part| part.get("text")?.as_str())
.collect::<Vec<_>>()
.join("\n"),
_ => return None,
};
Some((kind, text))
}
}
}
fn history_value_text(value: &Value) -> String {
value
.as_str()
.map_or_else(|| value.to_string(), str::to_owned)
}
fn history_chunk(text: &str, offset: usize, max_chars: usize) -> Result<(String, Option<usize>)> {
let start = text
.char_indices()
.map(|(index, _)| index)
.nth(offset)
.or_else(|| (text.chars().count() == offset).then_some(text.len()))
.ok_or_else(|| Error::Tool("history offset exceeds the item length".into()))?;
let mut chars = text[start..].chars();
let content = chars.by_ref().take(max_chars).collect::<String>();
let next = chars.next().is_some().then(|| offset + max_chars);
Ok((content, next))
}
fn history_excerpt(text: &str, query: &str) -> (usize, String) {
let lower = text.to_ascii_lowercase();
let query = query.to_ascii_lowercase();
let found = lower
.find(&query)
.or_else(|| query.split_whitespace().find_map(|term| lower.find(term)));
let offset = found.map_or(0, |index| text[..index].chars().count().saturating_sub(100));
(
offset,
text.chars()
.skip(offset)
.take(HISTORY_EXCERPT_CHARS)
.collect(),
)
}
async fn fork(context: MiddlewareCommandContext<'_>) -> Result<MiddlewareCommandOutput> {
if !context.arguments.trim().is_empty() {
return Ok(MiddlewareCommandOutput::render(
"sessions",
"! usage: fork",
FrontendTone::Warning,
));
}
let target = context.target;
let through_sequence = target
.as_ref()
.map_or(context.checkpoint.sequence, |target| {
target.checkpoint_sequence
});
let items = transcript_items_through(&context, through_sequence).await?;
let Some(target) = target else {
let options = fork_options(&items, context.session_id);
if options.is_empty() {
return Ok(MiddlewareCommandOutput::render(
"sessions",
"no messages to fork",
FrontendTone::Neutral,
));
}
return Ok(MiddlewareCommandOutput::events(vec![
FrontendEvent::Picker {
title: text::PICKER_FORK_CHAT_FROM_MESSAGE.into(),
options,
},
]));
};
let transcript = fork_prefix(items, &target, context.session_id)?;
let checkpoint = manual_fork_checkpoint(context.checkpoint, transcript);
context
.checkpoints
.fork(
&context.checkpoint.session_id,
target.checkpoint_sequence,
&checkpoint,
)
.await?;
Ok(MiddlewareCommandOutput::render(
"sessions",
format!("◇ forked chat {}", compact_id(&checkpoint.session_id)),
FrontendTone::Success,
))
}
async fn transcript_items_through(
context: &MiddlewareCommandContext<'_>,
through_sequence: u64,
) -> Result<Vec<(MessageTarget, serde_json::Value)>> {
let mut before_sequence =
Some(through_sequence.checked_add(1).ok_or_else(|| {
Error::Checkpoint("fork sequence exceeds the supported range".into())
})?);
let mut pages = Vec::new();
loop {
let page = context
.checkpoints
.transcript_page(
context.session_id,
TranscriptPageRequest {
before_sequence,
max_batches: DEFAULT_PAGE_SIZE,
},
)
.await?;
before_sequence = page.next_before_sequence;
pages.push(page.into_positioned_items_chronological());
if before_sequence.is_none() {
break;
}
}
Ok(pages.into_iter().rev().flatten().collect())
}
fn fork_options(
items: &[(MessageTarget, serde_json::Value)],
session_id: &str,
) -> Vec<FrontendPickerOption> {
replay_events(items, session_id)
.into_iter()
.rev()
.filter_map(|event| {
let (description, message, target) = match event {
EventMsg::Message(message) => (
text::PICKER_USER_MESSAGE,
message.text,
message.message_target?,
),
EventMsg::AssistantMessage(message) => (
text::PICKER_ASSISTANT_MESSAGE,
assistant_message_text(&message.content)?,
message.message_target?,
),
_ => return None,
};
Some(FrontendPickerOption {
label: compact_message(&message),
description: description.into(),
detail: String::new(),
symbol: None,
shows_detail: false,
op: Op::CapabilityCommand {
capability: MANIFEST.id.into(),
command: "fork".into(),
arguments: String::new(),
input: None,
target: Some(target),
},
})
})
.collect()
}
fn fork_prefix(
items: Vec<(MessageTarget, serde_json::Value)>,
target: &MessageTarget,
session_id: &str,
) -> Result<Vec<serde_json::Value>> {
let index = items
.iter()
.position(|(position, _)| position == target)
.ok_or_else(invalid_fork_target)?;
let prefix = &items[..=index];
if !replay_events(prefix, session_id)
.into_iter()
.any(|event| match event {
EventMsg::Message(message) => message.message_target.as_ref() == Some(target),
EventMsg::AssistantMessage(message) => message.message_target.as_ref() == Some(target),
_ => false,
})
{
return Err(invalid_fork_target());
}
Ok(items
.into_iter()
.take(index + 1)
.map(|(_, item)| item)
.collect())
}
fn invalid_fork_target() -> Error {
Error::Checkpoint("fork target is not a safe durable message boundary".into())
}
fn assistant_message_text(content: &[crate::protocol::ModelStepContent]) -> Option<String> {
[
crate::protocol::ModelStepContentPhase::FinalAnswer,
crate::protocol::ModelStepContentPhase::Commentary,
]
.into_iter()
.find_map(|phase| {
let text = content
.iter()
.filter(|item| item.phase == phase)
.map(|item| item.text.as_str())
.collect::<String>();
(!text.is_empty()).then_some(text)
})
}
fn compact_message(message: &str) -> String {
message
.split_whitespace()
.collect::<Vec<_>>()
.join(" ")
.chars()
.take(42)
.collect::<String>()
.trim_end()
.into()
}
fn manual_fork_checkpoint(parent: &Checkpoint, context: Vec<serde_json::Value>) -> Checkpoint {
let mut checkpoint = Checkpoint::empty(Uuid::new_v4().to_string());
checkpoint.context = context;
strip_attachment_references(&mut checkpoint.context);
checkpoint
.first_user_message
.clone_from(&parent.first_user_message);
checkpoint.model_route.clone_from(&parent.model_route);
checkpoint
.session_context
.clone_from(&parent.session_context);
checkpoint.metadata.clone_from(&parent.metadata);
checkpoint.session_context.origin_label = None;
checkpoint
}
async fn resume(
context: MiddlewareCommandContext<'_>,
page_size: usize,
) -> Result<MiddlewareCommandOutput> {
let arguments = context.arguments.trim();
let cursor = if arguments.is_empty() {
None
} else {
match serde_json::from_str(arguments) {
Ok(cursor) => Some(cursor),
Err(_) => {
return Ok(MiddlewareCommandOutput::render(
"sessions",
"! usage: resume",
FrontendTone::Warning,
));
}
}
};
let options = resume_options(&context, cursor, page_size).await?;
if options.is_empty() {
return Ok(MiddlewareCommandOutput::render(
"sessions",
"no saved chats",
FrontendTone::Neutral,
));
}
Ok(MiddlewareCommandOutput::events(vec![
FrontendEvent::Picker {
title: text::PICKER_RESUME_CHAT.into(),
options,
},
]))
}
fn compact_id(id: &str) -> &str {
id.get(..8).unwrap_or(id)
}
async fn resume_options(
context: &MiddlewareCommandContext<'_>,
cursor: Option<SessionCursor>,
page_size: usize,
) -> Result<Vec<FrontendPickerOption>> {
let page = context
.checkpoints
.list_sessions_page(SessionPageRequest {
bot_id: None,
cursor,
limit: page_size,
})
.await?;
resume_page_options(page, &context.checkpoint.session_id)
}
fn resume_page_options(
page: SessionPage,
current_session_id: &str,
) -> Result<Vec<FrontendPickerOption>> {
let mut options = page
.sessions
.into_iter()
.filter_map(|session| resume_option(session, current_session_id))
.collect::<Vec<_>>();
if let Some(cursor) = page.next_cursor {
options.push(FrontendPickerOption {
label: text::WIDGET_MORE_CHATS.into(),
description: String::new(),
detail: String::new(),
symbol: None,
shows_detail: false,
op: Op::CapabilityCommand {
capability: MANIFEST.id.into(),
command: "resume".into(),
arguments: serde_json::to_string(&cursor)?,
input: None,
target: None,
},
});
}
Ok(options)
}
fn resume_option(
session: SessionSummary,
current_session_id: &str,
) -> Option<FrontendPickerOption> {
if !session.catalog_visible || session.session_id == current_session_id {
return None;
}
let description = session_description(&session);
let label = session.first_user_message.map_or_else(
|| {
format!(
"{} {}",
if session.parent_session_id.is_some() {
"Fork"
} else {
"Chat"
},
compact_id(&session.session_id)
)
},
|message| compact_message(&message),
);
Some(FrontendPickerOption {
label,
description,
detail: String::new(),
symbol: None,
shows_detail: false,
op: Op::ResumeSession {
session_id: session.session_id,
},
})
}
fn session_description(session: &SessionSummary) -> String {
let mut details = [
session.session_context.workspace_label.as_deref(),
session.session_context.origin_label.as_deref(),
]
.into_iter()
.flatten()
.map(str::to_owned)
.collect::<Vec<_>>();
details.push(format!("created at Unix time {}", session.created_at));
details.join(" · ")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn history_tools_are_directly_available_for_the_required_bot() {
let state = tempfile::tempdir().expect("state");
let checkpoints: Arc<dyn crate::backend::checkpoint::CheckpointStore> = Arc::new(
crate::backend::checkpoint::sqlite::SqliteCheckpoint::new(
state.path().join("checkpoints.sqlite3"),
)
.expect("checkpoint store"),
);
let runtime = RuntimeContext {
sender: crate::agent::test_sender(),
checkpoints,
session_id: "session".into(),
model_route: "model".into(),
model: "model".into(),
approval_policy: crate::backend::sandbox::ApprovalPolicy::Ask,
session_context: crate::protocol::SessionContext {
bot_id: "bot-1".into(),
..crate::protocol::SessionContext::default()
},
metadata: std::collections::BTreeMap::new(),
role: crate::agent::AgentRole::Main,
frontend: Arc::new(|_| Ok(())),
};
let mut catalog = Catalog::default();
Sessions::default()
.register(&mut catalog, &runtime)
.expect("register session tools");
catalog.finalize().expect("finalize tools");
let names = catalog
.direct_definitions()
.iter()
.map(|tool| tool.name.clone())
.collect::<Vec<_>>();
assert_eq!(names, ["read_history", "search_history"]);
}
async fn save_history(
checkpoints: &dyn crate::backend::checkpoint::CheckpointStore,
session_id: &str,
bot_id: &str,
items: Vec<Value>,
) -> Checkpoint {
let mut checkpoint = Checkpoint::empty(session_id);
checkpoint.sequence = 1;
checkpoint.session_context.bot_id = bot_id.into();
checkpoint.context.clone_from(&items);
checkpoints
.save(&checkpoint, &items, None)
.await
.expect("save history");
checkpoint
}
fn history_store(path: &std::path::Path) -> History {
History {
checkpoints: Arc::new(
crate::backend::checkpoint::sqlite::SqliteCheckpoint::new(path).expect("store"),
),
session_id: "current".into(),
bot_id: "researcher".into(),
}
}
fn search_args(query: &str, scope: HistoryScope, cursor: Option<String>) -> SearchHistoryArgs {
SearchHistoryArgs {
query: query.into(),
scope,
cursor,
}
}
#[tokio::test]
async fn history_recovers_offloaded_tool_output_beyond_the_first_page() {
let state = tempfile::tempdir().expect("state");
let history = history_store(&state.path().join("history.sqlite3"));
let output = format!("{}\nneedle exact result\n", "🦀".repeat(9_000));
let mut checkpoint = save_history(history.checkpoints.as_ref(), "current", "researcher", vec![
crate::backend::model::user_message("Investigate"),
serde_json::json!({"type":"function_call", "call_id":"call-1", "name":"read_file", "arguments":"{\"path\":\"earlier.rs\"}"}),
crate::backend::model::tool_output("call-1", &output, false),
]).await;
checkpoint.context[2]["output"] = Value::String("[offloaded]".into());
for sequence in 2..=70 {
checkpoint.sequence = sequence;
history
.checkpoints
.save(
&checkpoint,
&[serde_json::json!({"role":"assistant", "content":"newer unrelated work"})],
None,
)
.await
.expect("save later batch");
}
let mut cursor = None;
let mut pages = 0;
let hit = loop {
let page: Value = serde_json::from_str(
&history
.search(search_args("needle", HistoryScope::Current, cursor))
.await
.expect("search"),
)
.expect("search page");
pages += 1;
if let Some(hit) = page["hits"].as_array().expect("hits").first() {
break hit.clone();
}
assert!(pages < 10, "history scan must make progress");
cursor = Some(
page["next_cursor"]
.as_str()
.expect("older history cursor")
.into(),
);
};
assert!(pages > 1);
assert_eq!(hit["session_id"], "current");
assert_eq!(hit["kind"], "tool_result");
assert!(
hit["excerpt"]
.as_str()
.expect("excerpt")
.contains("needle exact result")
);
let target: MessageTarget = serde_json::from_value(hit["target"].clone()).expect("target");
assert_eq!(
target,
MessageTarget {
checkpoint_sequence: 1,
batch_item_count: 3
}
);
let mut restored = String::new();
let mut offset = 0;
loop {
let page: Value = serde_json::from_str(
&history
.read(ReadHistoryArgs {
session_id: None,
target,
offset,
max_chars: 4_000,
})
.await
.expect("read history"),
)
.expect("read page");
restored.push_str(page["text"].as_str().expect("text"));
let Some(next) = page["next_offset"].as_u64() else {
break;
};
offset = usize::try_from(next).expect("offset");
}
assert_eq!(restored, output);
let calls: Value = serde_json::from_str(
&history
.read(ReadHistoryArgs {
session_id: None,
target: MessageTarget {
checkpoint_sequence: 1,
batch_item_count: 2,
},
offset: 0,
max_chars: 4_000,
})
.await
.expect("read call"),
)
.expect("call");
assert_eq!(calls["text"], "read_file\n{\"path\":\"earlier.rs\"}");
}
#[tokio::test]
async fn history_cursor_continues_inside_a_large_message() {
let state = tempfile::tempdir().expect("state");
let history = history_store(&state.path().join("history.sqlite3"));
save_history(
history.checkpoints.as_ref(),
"current",
"researcher",
vec![crate::backend::model::user_message(&format!(
"{}needle after the scan limit",
"padding ".repeat(10_000)
))],
)
.await;
let first: Value = serde_json::from_str(
&history
.search(search_args("needle", HistoryScope::Current, None))
.await
.expect("first page"),
)
.expect("first page json");
assert!(first["hits"].as_array().expect("hits").is_empty());
assert!(first["scanned_chars"].as_u64().expect("scan bound") <= 64_000);
let second: Value = serde_json::from_str(
&history
.search(search_args(
"needle",
HistoryScope::Current,
Some(first["next_cursor"].as_str().expect("cursor").into()),
))
.await
.expect("second page"),
)
.expect("second page json");
assert!(
second["hits"][0]["excerpt"]
.as_str()
.expect("late match")
.contains("needle after the scan limit")
);
}
#[tokio::test]
async fn history_requires_explicit_other_chat_scope_and_rejects_foreign_bot_reads() {
let state = tempfile::tempdir().expect("state");
let history = history_store(&state.path().join("history.sqlite3"));
for (session, bot) in [
("current", "researcher"),
("prior", "researcher"),
("private", "writer"),
] {
save_history(
history.checkpoints.as_ref(),
session,
bot,
vec![crate::backend::model::user_message(&format!(
"needle in {session}"
))],
)
.await;
}
for (scope, expected) in [
(HistoryScope::Current, "current"),
(HistoryScope::OtherChats, "prior"),
] {
let result: Value = serde_json::from_str(
&history
.search(search_args("needle", scope, None))
.await
.expect("search"),
)
.expect("result");
let hits = result["hits"].as_array().expect("hits");
assert_eq!(hits.len(), 1);
assert_eq!(hits[0]["session_id"], expected);
}
let target = MessageTarget {
checkpoint_sequence: 1,
batch_item_count: 1,
};
assert!(
history
.read(ReadHistoryArgs {
session_id: Some("prior".into()),
target,
offset: 0,
max_chars: 100
})
.await
.is_ok()
);
assert!(
history
.read(ReadHistoryArgs {
session_id: Some("private".into()),
target,
offset: 0,
max_chars: 100
})
.await
.is_err()
);
let mut forged = history
.cursor(search_args("needle", HistoryScope::OtherChats, None))
.expect("cursor");
forged.session_id = Some("private".into());
assert!(
history
.search(search_args(
"needle",
HistoryScope::OtherChats,
Some(serde_json::to_string(&forged).expect("cursor json"))
))
.await
.is_err()
);
forged.session_id = Some("prior".into());
assert!(
history
.search(search_args(
"needle",
HistoryScope::Current,
Some(serde_json::to_string(&forged).expect("cursor json"))
))
.await
.is_err()
);
}
#[tokio::test]
async fn history_validates_limits_targets_and_hides_private_context() {
let state = tempfile::tempdir().expect("state");
let history = history_store(&state.path().join("history.sqlite3"));
save_history(history.checkpoints.as_ref(), "current", "researcher", vec![
serde_json::json!({"role":"assistant", "content":"visible", "encrypted_content":"hidden", "_mobius_reasoning":"hidden"}),
serde_json::json!({"type":"reasoning", "encrypted_content":"secret"}),
crate::backend::model::internal_user_message("private", "hidden"),
]).await;
let target = MessageTarget {
checkpoint_sequence: 1,
batch_item_count: 1,
};
for (offset, max_chars) in [(0, 0), (0, MAX_HISTORY_READ_CHARS + 1), (usize::MAX, 10)] {
assert!(
history
.read(ReadHistoryArgs {
session_id: None,
target,
offset,
max_chars
})
.await
.is_err()
);
}
for batch_item_count in [0, 2, 3, 4] {
assert!(
history
.read(ReadHistoryArgs {
session_id: None,
target: MessageTarget {
batch_item_count,
..target
},
offset: 0,
max_chars: 100,
})
.await
.is_err()
);
}
assert!(
history
.cursor(search_args("", HistoryScope::Current, None))
.is_err()
);
assert!(
history
.cursor(search_args(
&"a".repeat(MAX_HISTORY_QUERY_BYTES + 1),
HistoryScope::Current,
None
))
.is_err()
);
assert!(
history
.cursor(search_args(
"needle",
HistoryScope::Current,
Some("x".repeat(MAX_HISTORY_CURSOR_BYTES + 1))
))
.is_err()
);
let result = history
.read(ReadHistoryArgs {
session_id: None,
target,
offset: 0,
max_chars: 100,
})
.await
.expect("visible text");
assert!(result.contains("visible"));
assert!(!result.contains("hidden"));
}
#[test]
fn sessions_rejects_page_sizes_outside_its_manifest_bounds() {
assert!(Sessions::new(0).is_err());
assert!(Sessions::new(MAX_PAGE_SIZE + 1).is_err());
}
#[test]
fn fork_is_exposed_as_a_generic_message_action() {
let contribution = Sessions::default().frontend();
let widget = contribution.widgets.first().expect("fork widget");
assert_eq!(widget.slot, FrontendSlot::MessageActions);
assert_eq!(widget.text, "Fork chat");
assert_eq!(widget.symbol, Some(FrontendSymbol::Branch));
assert_eq!(
widget.action,
Some(Op::CapabilityCommand {
capability: "sessions".into(),
command: "fork".into(),
arguments: String::new(),
input: None,
target: None,
})
);
}
#[test]
fn fork_picker_lists_only_safe_user_and_assistant_messages() {
let items = [
crate::backend::model::message_input(&crate::protocol::MessageEvent {
author: crate::protocol::MessageAuthor::User,
delivery: crate::protocol::MessageDelivery::Turn,
text: "Start here".into(),
attachments: Vec::new(),
reply: None,
message_target: None,
})
.expect("message input"),
serde_json::json!({"type": "function_call", "call_id": "call-1", "name": "read"}),
serde_json::json!({"type": "function_call_output", "call_id": "call-1", "output": "done"}),
serde_json::json!({"role": "assistant", "content": "Finished"}),
]
.into_iter()
.enumerate()
.map(|(index, item)| {
(
MessageTarget {
checkpoint_sequence: index as u64 + 1,
batch_item_count: 1,
},
item,
)
})
.collect::<Vec<_>>();
let targets = fork_options(&items, "session")
.into_iter()
.map(|option| match option.op {
Op::CapabilityCommand {
target: Some(target),
..
} => target,
operation => panic!("expected targeted fork, got {operation:?}"),
})
.collect::<Vec<_>>();
assert_eq!(targets, [items[3].0, items[0].0]);
}
#[test]
fn fork_prefix_includes_the_selected_batch_item() {
let items = vec![
(
MessageTarget {
checkpoint_sequence: 1,
batch_item_count: 1,
},
serde_json::json!({"role": "user", "content": "Question"}),
),
(
MessageTarget {
checkpoint_sequence: 2,
batch_item_count: 1,
},
serde_json::json!({"role": "assistant", "content": "Answer"}),
),
(
MessageTarget {
checkpoint_sequence: 2,
batch_item_count: 2,
},
serde_json::json!({"type": "function_call", "call_id": "call-1", "name": "read"}),
),
];
let target = items[1].0;
let prefix = fork_prefix(items, &target, "session").expect("fork prefix");
assert_eq!(
prefix,
[
serde_json::json!({"role": "user", "content": "Question"}),
serde_json::json!({"role": "assistant", "content": "Answer"}),
]
);
}
#[test]
fn fork_prefix_rejects_a_message_with_an_open_tool_call() {
let items = vec![
(
MessageTarget {
checkpoint_sequence: 1,
batch_item_count: 1,
},
serde_json::json!({"type": "function_call", "call_id": "call-1", "name": "read"}),
),
(
MessageTarget {
checkpoint_sequence: 1,
batch_item_count: 2,
},
serde_json::json!({"role": "assistant", "content": "Working"}),
),
];
let target = items[1].0;
let error = fork_prefix(items, &target, "session").expect_err("unsafe fork must fail");
assert_eq!(
error.to_string(),
"checkpoint error: fork target is not a safe durable message boundary"
);
}
#[test]
fn resume_lists_fresh_forks() {
let summary = |session_id: &str, parent_session_id: Option<&str>| SessionSummary {
session_id: session_id.into(),
session_context: Default::default(),
parent_session_id: parent_session_id.map(str::to_string),
parent_sequence: parent_session_id.map(|_| 4),
sequence: 0,
catalog_visible: true,
first_user_message: None,
execution_stats: Default::default(),
created_at: 0,
updated_at: 0,
};
assert_eq!(
resume_option(summary("branch-id", Some("parent")), "current")
.map(|option| option.label),
Some("Fork branch-i".into())
);
}
#[test]
fn resume_lists_empty_durable_root_chats() {
let option = resume_option(
SessionSummary {
session_id: "empty-root".into(),
session_context: Default::default(),
parent_session_id: None,
parent_sequence: None,
sequence: 0,
catalog_visible: true,
first_user_message: None,
execution_stats: Default::default(),
created_at: 0,
updated_at: 0,
},
"current",
);
assert_eq!(
option.map(|option| option.label),
Some("Chat empty-ro".into())
);
}
#[test]
fn resume_lists_catalog_visible_chats_across_workspaces() {
let summary = |session_id: &str, workspace: &str| SessionSummary {
session_id: session_id.into(),
session_context: crate::protocol::SessionContext {
workspace_id: Some(workspace.into()),
workspace_label: Some(workspace.into()),
..crate::protocol::SessionContext::default()
},
parent_session_id: None,
parent_sequence: None,
sequence: 1,
catalog_visible: true,
first_user_message: Some(format!("Work in {workspace}")),
execution_stats: Default::default(),
created_at: 0,
updated_at: 0,
};
let options = resume_page_options(
SessionPage {
sessions: vec![
summary("workspace-a-chat", "Workspace A"),
summary("workspace-b-chat", "Workspace B"),
],
next_cursor: None,
},
"current",
)
.expect("resume options");
let session_ids = options
.into_iter()
.map(|option| match option.op {
Op::ResumeSession { session_id } => session_id,
operation => panic!("expected resume operation, got {operation:?}"),
})
.collect::<Vec<_>>();
assert_eq!(session_ids, ["workspace-a-chat", "workspace-b-chat"]);
}
#[test]
fn resume_excludes_only_current_and_explicitly_hidden_chats() {
let summary = |session_id: &str, catalog_visible: bool| SessionSummary {
session_id: session_id.into(),
session_context: Default::default(),
parent_session_id: None,
parent_sequence: None,
sequence: 0,
catalog_visible,
first_user_message: None,
execution_stats: Default::default(),
created_at: 0,
updated_at: 0,
};
let options = resume_page_options(
SessionPage {
sessions: vec![
summary("current", true),
summary("hidden", false),
summary("visible", true),
],
next_cursor: None,
},
"current",
)
.expect("resume options");
let session_ids = options
.into_iter()
.map(|option| match option.op {
Op::ResumeSession { session_id } => session_id,
operation => panic!("expected resume operation, got {operation:?}"),
})
.collect::<Vec<_>>();
assert_eq!(session_ids, ["visible"]);
}
#[test]
fn resume_description_includes_workspace_and_origin_labels() {
let option = resume_option(
SessionSummary {
session_id: "routine".into(),
session_context: crate::protocol::SessionContext {
workspace_label: Some("Project One".into()),
origin_label: Some("routine".into()),
..crate::protocol::SessionContext::default()
},
parent_session_id: None,
parent_sequence: None,
sequence: 1,
catalog_visible: true,
first_user_message: Some("Update dependencies".into()),
execution_stats: Default::default(),
created_at: 42,
updated_at: 42,
},
"current",
)
.expect("resume option");
assert_eq!(
option.description,
"Project One · routine · created at Unix time 42"
);
}
#[test]
fn manual_fork_keeps_context_workspace_and_metadata_but_clears_origin() {
let mut parent = Checkpoint::empty("parent");
parent.context = vec![serde_json::json!({
"role": "user",
"content": "Hello",
"_mobius_attachments": [{
"id": "378b8581-e96c-4413-a138-93e74561cb87",
"name": "photo.png",
"size": 1,
"media_type": "image/png"
}]
})];
parent.first_user_message = Some("Hello".into());
parent.metadata.insert(
"gateway.chat".into(),
serde_json::json!({"workspace": "/srv/project"}),
);
parent.session_context = crate::protocol::SessionContext {
bot_id: "bot-1".into(),
workspace_id: Some("workspace-1".into()),
workspace_label: Some("Project One".into()),
origin_label: Some("routine".into()),
..crate::protocol::SessionContext::default()
};
let fork = manual_fork_checkpoint(&parent, parent.context.clone());
assert!(fork.context[0].get("_mobius_attachments").is_none());
assert_eq!(fork.first_user_message, parent.first_user_message);
assert_eq!(fork.metadata, parent.metadata);
assert_eq!(
fork.session_context,
crate::protocol::SessionContext {
bot_id: "bot-1".into(),
workspace_id: Some("workspace-1".into()),
workspace_label: Some("Project One".into()),
..crate::protocol::SessionContext::default()
}
);
}
#[test]
fn resume_page_preserves_the_next_catalog_cursor() {
let cursor = SessionCursor {
updated_at: 12,
sequence: 4,
session_id: "next".into(),
};
let options = resume_page_options(
SessionPage {
sessions: Vec::new(),
next_cursor: Some(cursor.clone()),
},
"current",
)
.expect("build resume page");
let Op::CapabilityCommand { arguments, .. } = &options[0].op else {
panic!("expected middleware command");
};
assert_eq!(
serde_json::from_str::<SessionCursor>(arguments).expect("decode cursor"),
cursor
);
}
}