use std::sync::Arc; use async_trait::async_trait; use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use schemars::JsonSchema; use serde::Deserialize; use crate::feature::{ FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ToolContribution, ToolDeclaration, }; use crate::session_capture::{ ReadDetail, ReadOptions, ReadSelector, ReferenceKind, SearchOptions, SessionCapture, SessionEntryRef, ToolPart, }; const SHOW_OVERVIEW_DESCRIPTION: &str = "Show a sparse, bounded index of real user and assistant session entries."; const SEARCH_ENTRIES_DESCRIPTION: &str = "Search or compactly list a bounded range of committed session entries. Reasoning entries are never exposed."; const READ_ENTRY_DESCRIPTION: &str = "Read one committed session entry by stable SessionEntryRef. Compact mode is the default; full mode includes bounded tool input or output."; const DEFAULT_PAGE_LIMIT: usize = 20; const MAX_PAGE_LIMIT: usize = 100; const DEFAULT_READ_ITEMS: usize = 10; const MAX_READ_ITEMS: usize = 50; const MAX_READ_BYTES: usize = 16 * 1024; #[derive(Clone)] pub(crate) struct SessionExploreState { view: Arc, } impl SessionExploreState { pub(crate) fn new(view: SessionCapture) -> Self { Self { view: Arc::new(view), } } pub(crate) fn view(&self) -> &SessionCapture { &self.view } } #[derive(Clone)] pub(crate) struct SessionExploreFeature { state: SessionExploreState, } impl SessionExploreFeature { pub(crate) fn new(state: SessionExploreState) -> Self { Self { state } } } impl FeatureModule for SessionExploreFeature { fn descriptor(&self) -> FeatureDescriptor { FeatureDescriptor::builtin("session-explore", "Session Explore") .with_description( "Read-only exploration of one immutable host-provided session capture.", ) .with_tool(ToolDeclaration::new( "ShowOverview", SHOW_OVERVIEW_DESCRIPTION, )) .with_tool(ToolDeclaration::new( "SearchEntries", SEARCH_ENTRIES_DESCRIPTION, )) .with_tool(ToolDeclaration::new("ReadEntry", READ_ENTRY_DESCRIPTION)) } fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { context.tools().register(ToolContribution::new( "ShowOverview", show_overview_definition(self.state.clone()), ))?; context.tools().register(ToolContribution::new( "SearchEntries", search_entries_definition(self.state.clone()), ))?; context.tools().register(ToolContribution::new( "ReadEntry", read_entry_definition(self.state.clone()), ))?; Ok(()) } } fn show_overview_definition(state: SessionExploreState) -> ToolDefinition { Arc::new(move || { let schema = serde_json::to_value(schemars::schema_for!(ShowOverviewParams)) .unwrap_or_else(|_| serde_json::json!({})); let meta = ToolMeta::new("ShowOverview") .description(SHOW_OVERVIEW_DESCRIPTION) .input_schema(schema); let tool: Arc = Arc::new(ShowOverviewTool { state: state.clone(), }); (meta, tool) }) } fn search_entries_definition(state: SessionExploreState) -> ToolDefinition { Arc::new(move || { let schema = serde_json::to_value(schemars::schema_for!(SearchEntriesParams)) .unwrap_or_else(|_| serde_json::json!({})); let meta = ToolMeta::new("SearchEntries") .description(SEARCH_ENTRIES_DESCRIPTION) .input_schema(schema); let tool: Arc = Arc::new(SearchEntriesTool { state: state.clone(), }); (meta, tool) }) } fn read_entry_definition(state: SessionExploreState) -> ToolDefinition { Arc::new(move || { let schema = serde_json::to_value(schemars::schema_for!(ReadEntryParams)) .unwrap_or_else(|_| serde_json::json!({})); let meta = ToolMeta::new("ReadEntry") .description(READ_ENTRY_DESCRIPTION) .input_schema(schema); let tool: Arc = Arc::new(ReadEntryTool { state: state.clone(), }); (meta, tool) }) } #[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct ShowOverviewParams { #[serde(default)] offset: usize, #[serde(default)] limit: Option, } #[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct SearchEntriesParams { #[serde(default)] query: String, #[serde(default)] kind: Option, #[serde(default)] tool_part: Option, #[serde(default)] tool_name: Option, #[serde(default)] from: Option, #[serde(default)] through: Option, #[serde(default)] offset: usize, #[serde(default)] limit: Option, } #[derive(Debug, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] struct ReadEntryParams { entry_ref: String, #[serde(default = "default_read_mode")] mode: String, #[serde(default)] max_items: Option, } fn default_read_mode() -> String { "compact".to_string() } struct ShowOverviewTool { state: SessionExploreState, } #[async_trait] impl Tool for ShowOverviewTool { async fn execute( &self, input_json: &str, _context: llm_engine::tool::ToolExecutionContext, ) -> Result { let params: ShowOverviewParams = parse_input("ShowOverview", input_json)?; let limit = bounded_limit(params.limit, DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT); let overview = self.state.view().overview(); let page = overview .iter() .skip(params.offset) .take(limit) .map(|entry| { serde_json::json!({ "entry_ref": entry.id, "entry_range": entry.entry_range, "kind": entry.kind.as_str(), "label": entry.label, "text": entry.text, "intervening_entries": entry.intervening_entries, }) }) .collect::>(); let has_more = params.offset.saturating_add(page.len()) < overview.len(); json_output( format!("Showing {} session overview entrie(s).", page.len()), serde_json::json!({ "entries": page, "offset": params.offset, "next_offset": has_more.then_some(params.offset + page.len()), "total": overview.len(), }), ) } } struct SearchEntriesTool { state: SessionExploreState, } #[async_trait] impl Tool for SearchEntriesTool { async fn execute( &self, input_json: &str, _context: llm_engine::tool::ToolExecutionContext, ) -> Result { let params: SearchEntriesParams = parse_input("SearchEntries", input_json)?; let kind = params.kind.as_deref().map(parse_kind).transpose()?; let tool_part = params .tool_part .as_deref() .map(parse_tool_part) .transpose()?; let from = params.from.as_deref().map(parse_entry_ref).transpose()?; let through = params.through.as_deref().map(parse_entry_ref).transpose()?; if let (Some(from), Some(through)) = (&from, &through) { if from.source_index() > through.source_index() { return Err(ToolError::InvalidArgument( "SearchEntries from must not be after through".to_string(), )); } } let limit = bounded_limit(params.limit, DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT); let hits = self.state.view().search(&SearchOptions { query: params.query, kind, tool_part, tool_name: params.tool_name, limit: Some(limit), min_entry_index: None, from, through, offset: params.offset, }); let entries = hits .iter() .map(|hit| { serde_json::json!({ "entry_ref": hit.id, "kind": hit.kind.as_str(), "tool_part": hit.tool_part.map(|part| format!("{part:?}").to_lowercase()), "tool_name": hit.tool_name, "label": hit.label, "text": hit.summary, }) }) .collect::>(); json_output( format!("Found {} session entrie(s).", entries.len()), serde_json::json!({ "entries": entries, "offset": params.offset, "next_offset": (entries.len() == limit).then_some(params.offset + entries.len()), }), ) } } struct ReadEntryTool { state: SessionExploreState, } #[async_trait] impl Tool for ReadEntryTool { async fn execute( &self, input_json: &str, _context: llm_engine::tool::ToolExecutionContext, ) -> Result { let params: ReadEntryParams = parse_input("ReadEntry", input_json)?; let entry_ref = parse_entry_ref(¶ms.entry_ref)?; let detail = match params.mode.as_str() { "compact" => ReadDetail::Compact, "full" => ReadDetail::Full, other => { return Err(ToolError::InvalidArgument(format!( "invalid ReadEntry mode {other:?}; expected compact or full" ))); } }; let read = self.state.view().read( ReadSelector::Id(entry_ref.as_str()), ReadOptions { include_tools: true, tool_part: ToolPart::Both, detail, max_items: params .max_items .unwrap_or(DEFAULT_READ_ITEMS) .clamp(1, MAX_READ_ITEMS), max_bytes: MAX_READ_BYTES, }, ); let entries = read .entries .iter() .map(|entry| { serde_json::json!({ "entry_ref": entry.id, "entry_range": entry.entry_range, "kind": entry.kind.as_str(), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), "tool_name": entry.tool_name, "label": entry.label, "text": entry.text, }) }) .collect::>(); if entries.is_empty() { return Err(ToolError::ExecutionFailed( "session entry was not found in the host-provided capture".to_string(), )); } json_output( format!("Read {} session entrie(s).", entries.len()), serde_json::json!({ "entries": entries, "truncated": read.truncated, }), ) } } fn parse_input( tool_name: &str, input_json: &str, ) -> Result { serde_json::from_str(input_json) .map_err(|error| ToolError::InvalidArgument(format!("invalid {tool_name} input: {error}"))) } fn parse_entry_ref(value: &str) -> Result { SessionEntryRef::parse(value).ok_or_else(|| { ToolError::InvalidArgument(format!( "invalid SessionEntryRef {value:?}; expected E followed by a decimal source index" )) }) } fn parse_kind(value: &str) -> Result { ReferenceKind::parse(value).ok_or_else(|| { ToolError::InvalidArgument(format!( "invalid kind {value:?}; expected user, assistant, or tool" )) }) } fn parse_tool_part(value: &str) -> Result { ToolPart::parse(value).ok_or_else(|| { ToolError::InvalidArgument(format!( "invalid tool_part {value:?}; expected input, output, or both" )) }) } fn bounded_limit(requested: Option, default: usize, maximum: usize) -> usize { requested.unwrap_or(default).clamp(1, maximum) } fn json_output(summary: String, value: serde_json::Value) -> Result { let content = serde_json::to_string_pretty(&value) .map_err(|error| ToolError::ExecutionFailed(format!("serialize tool output: {error}")))?; Ok(ToolOutput { summary, content: Some(content), }) } #[cfg(test)] mod tests { use llm_engine::Item; use crate::feature::{FeatureRegistryBuilder, HookRegistryBuilder}; use super::*; fn state() -> SessionExploreState { SessionExploreState::new(SessionCapture::new( "segment-1", vec![ Item::user_message("first"), Item::reasoning("private"), Item::tool_call("call-1", "ExampleTool", "{}"), Item::assistant_message("second"), ], )) } #[test] fn session_explore_installs_read_only_tools_without_memory_extract() { let mut pending_tools = Vec::new(); let mut hook_builder = HookRegistryBuilder::default(); let report = FeatureRegistryBuilder::new() .with_module(SessionExploreFeature::new(state())) .install_into_pending(&mut pending_tools, &mut hook_builder); assert!(report.reports[0].installed); assert_eq!( report.installed_tool_names(), ["ShowOverview", "SearchEntries", "ReadEntry"] ); } #[tokio::test] async fn overview_uses_real_entry_refs_and_rejects_unknown_fields() { let tool = show_overview_definition(state())().1; let output = tool .execute("{}", llm_engine::tool::ToolExecutionContext::direct()) .await .unwrap(); let content = output.content.unwrap(); assert!(content.contains("E00000000")); assert!(content.contains("E00000003")); assert!(content.contains("\"intervening_entries\": 1")); assert!(!content.contains("private")); let error = tool .execute( r#"{"unexpected":true}"#, llm_engine::tool::ToolExecutionContext::direct(), ) .await .unwrap_err(); assert!(format!("{error:?}").contains("unknown field")); } #[tokio::test] async fn search_range_and_read_share_session_entry_refs() { let search = search_entries_definition(state())().1; let output = search .execute( r#"{"from":"E00000003","through":"E00000003"}"#, llm_engine::tool::ToolExecutionContext::direct(), ) .await .unwrap(); let content = output.content.unwrap(); assert!(content.contains("E00000003")); assert!(!content.contains("E00000000")); let read = read_entry_definition(state())().1; let output = read .execute( r#"{"entry_ref":"E00000003"}"#, llm_engine::tool::ToolExecutionContext::direct(), ) .await .unwrap(); assert!(output.content.unwrap().contains("second")); } }