feat: add compact session exploration tools

This commit is contained in:
Keisuke Hirata 2026-05-28 12:31:44 +09:00
parent c274e4a891
commit 12dd35cfb2
4 changed files with 345 additions and 7 deletions

View File

@ -84,6 +84,39 @@ struct SummaryParams {
pub text: String,
}
/// Input to `search_session_log`.
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct SearchSessionParams {
/// Case-insensitive substring to search in compact-target history.
pub query: String,
/// 0-based item offset to start searching from.
#[serde(default)]
pub offset: Option<usize>,
/// Maximum number of hits to return.
#[serde(default)]
pub limit: Option<usize>,
}
/// Input to `read_session_items`.
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct ReadSessionParams {
/// 0-based compact-target history item offset.
pub offset: usize,
/// Maximum number of items to return.
pub limit: usize,
/// `compact` omits tool arguments/full results; `full` includes message text and tool result content.
#[serde(default = "default_session_read_mode")]
pub mode: String,
}
fn default_session_read_mode() -> String {
"compact".to_string()
}
const SESSION_TOOL_MAX_OUTPUT_TOKENS: u64 = 12_000;
const SESSION_SEARCH_MAX_RESULTS: usize = 50;
const SESSION_READ_MAX_ITEMS: usize = 80;
const MARK_DESCRIPTION: &str = "Inject a file's contents into the compacted context so the \
next session starts with it already read. Use this for files the next task needs in full. \
Optionally specify `offset` (0-based line) and `limit` (line count) to inject only a slice. \
@ -98,6 +131,236 @@ const SUMMARY_DESCRIPTION: &str = "Provide the final structured summary text. Su
replace the previous content; only the last call is used. Must be called before the compact run \
ends or compaction fails.";
const SEARCH_SESSION_DESCRIPTION: &str = "Search the compact-target session history by \
case-insensitive substring. Returns item indexes and compact snippets. Use this when the initial \
overview is not enough to identify which part of the session matters. Results are bounded; narrow \
the query if important details are omitted.";
const READ_SESSION_DESCRIPTION: &str = "Read a bounded range of compact-target session history \
items by 0-based index. mode='compact' omits tool arguments, full tool results, and reasoning \
bodies; mode='full' includes message text and tool result content but still remains bounded. Use \
this to verify details before writing the summary.";
struct SessionLogToolState {
items: Arc<Vec<Item>>,
}
struct SearchSessionLogTool {
state: Arc<SessionLogToolState>,
}
#[async_trait]
impl Tool for SearchSessionLogTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
let params: SearchSessionParams = serde_json::from_str(input_json).map_err(|e| {
ToolError::InvalidArgument(format!("invalid search_session_log input: {e}"))
})?;
let query = params.query.trim().to_lowercase();
if query.is_empty() {
return Err(ToolError::InvalidArgument(
"search_session_log query must not be empty".to_string(),
));
}
let offset = params.offset.unwrap_or(0).min(self.state.items.len());
let limit = params
.limit
.unwrap_or(20)
.clamp(1, SESSION_SEARCH_MAX_RESULTS);
let mut hits = Vec::new();
for (idx, item) in self.state.items.iter().enumerate().skip(offset) {
let haystack = session_item_search_text(item).to_lowercase();
if haystack.contains(&query) {
hits.push(format_session_item(
idx,
item,
SessionReadMode::Compact,
600,
));
if hits.len() >= limit {
break;
}
}
}
let mut content = hits.join("\n\n");
let truncated = truncate_to_token_budget(&mut content, SESSION_TOOL_MAX_OUTPUT_TOKENS);
let summary = if hits.is_empty() {
format!("No session log hits for {query:?} from item offset {offset}.")
} else if truncated {
format!(
"Found {} session log hit(s) for {query:?}; output truncated. Narrow the query.",
hits.len()
)
} else {
format!("Found {} session log hit(s) for {query:?}.", hits.len())
};
Ok(ToolOutput {
summary,
content: (!content.is_empty()).then_some(content),
})
}
}
struct ReadSessionItemsTool {
state: Arc<SessionLogToolState>,
}
#[async_trait]
impl Tool for ReadSessionItemsTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
let params: ReadSessionParams = serde_json::from_str(input_json).map_err(|e| {
ToolError::InvalidArgument(format!("invalid read_session_items input: {e}"))
})?;
let mode = SessionReadMode::parse(&params.mode)?;
let offset = params.offset.min(self.state.items.len());
let limit = params.limit.clamp(1, SESSION_READ_MAX_ITEMS);
let end = offset.saturating_add(limit).min(self.state.items.len());
let mut blocks = Vec::new();
for idx in offset..end {
blocks.push(format_session_item(
idx,
&self.state.items[idx],
mode,
4_000,
));
}
let mut content = blocks.join("\n\n");
let truncated = truncate_to_token_budget(&mut content, SESSION_TOOL_MAX_OUTPUT_TOKENS);
let summary = if truncated {
format!(
"Read session items {offset}..{end} in {mode:?} mode; output truncated. Narrow the range."
)
} else {
format!("Read session items {offset}..{end} in {mode:?} mode.")
};
Ok(ToolOutput {
summary,
content: (!content.is_empty()).then_some(content),
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SessionReadMode {
Compact,
Full,
}
impl SessionReadMode {
fn parse(value: &str) -> Result<Self, ToolError> {
match value {
"compact" => Ok(Self::Compact),
"full" => Ok(Self::Full),
other => Err(ToolError::InvalidArgument(format!(
"invalid read_session_items mode {other:?}; expected 'compact' or 'full'"
))),
}
}
}
fn session_item_search_text(item: &Item) -> String {
match item {
Item::Message { role, content, .. } => format!(
"{:?} {}",
role,
content
.iter()
.map(|p| p.as_text())
.collect::<Vec<_>>()
.join("")
),
Item::ToolCall {
name, arguments, ..
} => format!("tool_call {name} {arguments}"),
Item::ToolResult {
summary, content, ..
} => format!(
"tool_result {summary} {}",
content.as_deref().unwrap_or_default()
),
Item::Reasoning { text, summary, .. } => format!("reasoning {text} {}", summary.join(" ")),
}
}
fn format_session_item(idx: usize, item: &Item, mode: SessionReadMode, max_chars: usize) -> String {
match item {
Item::Message { role, content, .. } => {
let text = content
.iter()
.map(|p| p.as_text())
.collect::<Vec<_>>()
.join("");
format!(
"[{idx} Message {:?}] {}",
role,
truncate_chars(&text, max_chars)
)
}
Item::ToolCall {
name, arguments, ..
} => match mode {
SessionReadMode::Compact => format!("[{idx} ToolCall] {name} (arguments omitted)"),
SessionReadMode::Full => format!(
"[{idx} ToolCall] {name}\narguments: {}",
truncate_chars(arguments, max_chars)
),
},
Item::ToolResult {
summary,
content,
is_error,
..
} => match mode {
SessionReadMode::Compact => format!(
"[{idx} ToolResult{}] {} (content omitted)",
if *is_error { " error" } else { "" },
truncate_chars(summary, 800)
),
SessionReadMode::Full => format!(
"[{idx} ToolResult{}] {}\ncontent: {}",
if *is_error { " error" } else { "" },
truncate_chars(summary, 800),
truncate_chars(content.as_deref().unwrap_or(""), max_chars)
),
},
Item::Reasoning { summary, .. } => match mode {
SessionReadMode::Compact => format!(
"[{idx} Reasoning] {} (body omitted)",
truncate_chars(&summary.join(" "), 800)
),
SessionReadMode::Full => format!(
"[{idx} Reasoning] {} (body omitted)",
truncate_chars(&summary.join(" "), 800)
),
},
}
}
fn truncate_chars(text: &str, max_chars: usize) -> String {
if text.chars().count() <= max_chars {
return text.to_string();
}
let mut out = text.chars().take(max_chars).collect::<String>();
out.push_str("… [truncated]");
out
}
fn truncate_to_token_budget(text: &mut String, max_tokens: u64) -> bool {
let max_bytes = max_tokens.saturating_mul(4) as usize;
if text.len() <= max_bytes {
return false;
}
let mut cut = 0;
for (idx, _) in text.char_indices() {
if idx > max_bytes {
break;
}
cut = idx;
}
text.truncate(cut);
text.push_str("\n… [session tool output truncated]");
true
}
struct MarkReadRequiredTool {
fs: ScopedFs,
ctx: Arc<Mutex<CompactWorkerContext>>,
@ -247,6 +510,36 @@ pub(crate) fn write_summary_tool(ctx: Arc<Mutex<CompactWorkerContext>>) -> ToolD
})
}
pub(crate) fn search_session_log_tool(items: Arc<Vec<Item>>) -> ToolDefinition {
let state = Arc::new(SessionLogToolState { items });
Arc::new(move || {
let schema = schemars::schema_for!(SearchSessionParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("search_session_log")
.description(SEARCH_SESSION_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(SearchSessionLogTool {
state: state.clone(),
});
(meta, tool)
})
}
pub(crate) fn read_session_items_tool(items: Arc<Vec<Item>>) -> ToolDefinition {
let state = Arc::new(SessionLogToolState { items });
Arc::new(move || {
let schema = schemars::schema_for!(ReadSessionParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("read_session_items")
.description(READ_SESSION_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(ReadSessionItemsTool {
state: state.clone(),
});
(meta, tool)
})
}
/// Interceptor that monitors compact-worker context occupancy.
///
/// `max_input_tokens` remains the hard circuit breaker. Before that point,
@ -516,6 +809,45 @@ mod tests {
assert_eq!(guard.references[0], PathBuf::from(p));
}
#[tokio::test]
async fn search_session_log_returns_bounded_hits_without_full_tool_content() {
let items = Arc::new(vec![
Item::user_message("investigate compact failure"),
Item::tool_result_with_content(
"call-1",
"read trace with compact failure",
"very large raw trace body with secret detail",
),
]);
let tool: Arc<dyn Tool> = Arc::new(SearchSessionLogTool {
state: Arc::new(SessionLogToolState { items }),
});
let input = serde_json::json!({ "query": "compact", "limit": 10 }).to_string();
let out = tool.execute(&input).await.unwrap();
let content = out.content.unwrap();
assert!(content.contains("investigate compact failure"));
assert!(content.contains("read trace with compact failure"));
assert!(!content.contains("secret detail"));
}
#[tokio::test]
async fn read_session_items_full_mode_can_read_tool_result_content() {
let items = Arc::new(vec![Item::tool_result_with_content(
"call-1",
"read trace",
"raw trace detail",
)]);
let tool: Arc<dyn Tool> = Arc::new(ReadSessionItemsTool {
state: Arc::new(SessionLogToolState { items }),
});
let input = serde_json::json!({ "offset": 0, "limit": 1, "mode": "full" }).to_string();
let out = tool.execute(&input).await.unwrap();
let content = out.content.unwrap();
assert!(content.contains("raw trace detail"));
}
#[test]
fn slice_lines_handles_offset_and_limit() {
let text = "a\nb\nc\nd";

View File

@ -2296,7 +2296,8 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
pub async fn compact(&mut self, retained_tokens: u64) -> Result<SegmentId, PodError> {
use crate::compact::worker::{
CompactWorkerContext, CompactWorkerInterceptor, add_reference_tool,
mark_read_required_tool, write_summary_tool,
mark_read_required_tool, read_session_items_tool, search_session_log_tool,
write_summary_tool,
};
use crate::fs_view::PodFsView;
@ -2447,9 +2448,12 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
));
summary_worker.set_max_turns(worker_max_turns);
// Tools: read_file (shared scope, fresh tracker) + the three
// compact-specific tools that populate `ctx`.
// Tools: read_file (shared scope, fresh tracker), bounded session
// history exploration, and compact-specific tools that populate `ctx`.
let compact_target_items = Arc::new(items_to_summarise.clone());
summary_worker.register_tool(tools::read_tool(scoped_fs.clone(), summary_tracker));
summary_worker.register_tool(search_session_log_tool(compact_target_items.clone()));
summary_worker.register_tool(read_session_items_tool(compact_target_items));
summary_worker.register_tool(mark_read_required_tool(scoped_fs.clone(), ctx.clone()));
summary_worker.register_tool(add_reference_tool(ctx.clone()));
summary_worker.register_tool(write_summary_tool(ctx.clone()));

View File

@ -211,6 +211,7 @@ write_summary(text) — 構造化要約を出力/上書き
- デフォルトリファレンスの一覧
- TaskStore snapshot
3. compact worker が自律的に:
- search_session_log / read_session_items で bounded overview から漏れた compact 対象履歴を必要範囲だけ探索
- read_file で各ファイルを読み、必要性を判断
- mark_read_required / add_reference で指定
- write_summary で構造化要約を出力(呼び直し可)

View File

@ -5,10 +5,11 @@ The conversation input is a bounded overview/index, not the full transcript. Tre
## Workflow
1. Read the provided overview/index and current TaskStore snapshot.
2. Use `read_file` to inspect referenced files before deciding what the next session needs. Prefer skimming over blind inclusion.
3. For files whose current contents are load-bearing for the active work, call `mark_read_required` to inject them into the next session. These count against the auto-read token budget — spend it deliberately.
4. For files the next session should know about but can fetch on demand, call `add_reference` to record the path without embedding contents.
5. Finish with `write_summary` carrying the final text. You may call it multiple times; only the last call is kept.
2. If the overview does not contain enough detail, use `search_session_log` to find relevant compact-target history items, then `read_session_items` to inspect only the needed range.
3. Use `read_file` to inspect referenced files before deciding what the next session needs. Prefer skimming over blind inclusion.
4. For files whose current contents are load-bearing for the active work, call `mark_read_required` to inject them into the next session. These count against the auto-read token budget — spend it deliberately.
5. For files the next session should know about but can fetch on demand, call `add_reference` to record the path without embedding contents.
6. Finish with `write_summary` carrying the final text. You may call it multiple times; only the last call is kept.
Stop nominating and close out with `write_summary` as soon as the auto-read budget is exhausted, when a compact worker budget warning arrives, or whenever further exploration would not change the next session's next step.